sonogram.py

import sonogram_utility as su
from pyannote.audio import Pipeline
import pickle
import torch
import soundfile as sf
import numpy as np
from pyannote.core import Segment
from pyannote.audio.models.segmentation import PyanNet
from pyannote.audio import Inference
from pyannote.pipeline.parameter import ParamDict
import torch

class Sonogram():
    '''
    A class to hold the Sonogram model

    ...

    Attributes
    ----------
    earlyCleanup : bool
        Determines whether temporary arrays should be deleted ASAP
    isTPU : bool
        Determines whether TPU has been detected and utilized
    isGPU : bool
        Determines whether GPU has been detected and utilized
    device : torch.device
        Device to use for accelerated processing
    version : str
        Named version to determine which Sonogram model to load
    pipeline : pyannote.audio.Pipeline
        Representation of model as pipeline via pyannote
    groupClassifier : sklearn.svm.SVC
        SVM reclassifier

    Methods
    -------
    classifyEmbedding(embedding)
        Classifies 10 second feature embedding using groupClassifier
    processFile(filePath)
        Loads and processes file to provide diarization and analysis context
    activeSpeaker(inAnnotation,step=1)
        Determines the single active speaker for each timestep
    annotationToNoiseList(inAnnotation,maxTime,stepSize=2,windowSize=90)
        Determines which noise category applies to each timestep
    toDevice()
        Moves pipeline to device
    toCPU()
        Moves pipeline to CPU
    '''
    
    def __init__(self,version='1.0'):
        '''
        Parameters
        ----------
        version : str
            The named version of Sonogram to load
        '''
        self.earlyCleanup = True
        
        self.isTPU = False
        self.isGPU = False
        # Check if TPU or GPU are available
        try:
            # Force expected error as TPU has not yet been validated or necessary
            raise(RuntimeError("Not an error"))
            #device = xm.xla_device()
            print("TPU is available.")
            self.isTPU = True
        except RuntimeError as e:
            print(f"TPU is not available: {e}")
            self.isGPU = torch.cuda.is_available()
            if not self.isGPU:
                print(f"GPU is not available")
            # Fallback to CPU or other devices if needed
            self.device = torch.device("cuda" if self.isGPU else "cpu")
            print(f"Using {self.device} instead.")

        self.version = version
        # pyannote pre-trained version
        if version == 'speaker-diarization-3.1':
            self.pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization-3.1")
        # Sonogram trained version as of 20251208
        elif version == '1.0':
            baselinePipeline = Pipeline.from_pretrained("pyannote/speaker-diarization-3.1")
            newSpecs = baselinePipeline._segmentation.model.specifications
            segModel = PyanNet.from_pretrained('20251208_Sonogram_Segmentation.ckpt')
            segModel.specifications = newSpecs
            
            segmentation_duration = segModel.specifications.duration
            baselinePipeline._segmentation = Inference(
                segModel,
                duration=segmentation_duration,
                step=baselinePipeline.segmentation_step * segmentation_duration,
                skip_aggregation=True,
                batch_size=1,
            )
            
            baselinePipeline.segmentation = ParamDict(min_duration_off=0.0,)
            self.pipeline = baselinePipeline

        # Load SVM reclassifier
        with open('05062026_groupClassifier.pkl', 'rb') as f:
            self.groupClassifier = pickle.load(f)

    def classifyEmbedding(self,embedding):
        '''
        Classifies 10 second feature embedding using groupClassifier

        ...
        
        Parameters
        ----------
        embedding : np.array(x,x)
            10 second feature embedding to classify

        Returns
        -------
        _ : int
            0 for No voice, 1 for Individual voice, 2 for Indistinguishable Group voices
        '''
        return int(self.groupClassifier.predict(embedding.reshape(1, -1)).item())
    
    def processFile(self,filePath):
        '''
        Loads and processes file to provide diarization and analysis context

        ...
        
        Parameters
        ----------
        filePath : str
            Full path to audio file to process

        Returns
        -------
        diarization : pyannote.core.Annotation
            Diarization result from pipeline
        totalTimeInSeconds : int
            Approximate length of audio for use in diagrams
        waveform : np.array
            Audio waveform as loaded from file
        sample_rate : int
            Sample rate of loaded audio
        '''
        # Loading audio file
        print(f"Loading file: {filePath}")
        data, sample_rate = sf.read(filePath, dtype="float32", always_2d=True)
        waveform = torch.from_numpy(data.T)  # shape: [channels, samples]
        # Wrapping as AudioFile
        audioFile = {"waveform": waveform, "sample_rate": sample_rate}

        # Much of following code is modified from 
        #        pyannote.audio.pipelines.speaker_diarization.SpeakerDiarization
        print("Detecting Voices")
        segmentations = self.pipeline.get_segmentations(audioFile)
        print("Generating vocal embeddings")
        embeddings = self.pipeline.get_embeddings(audioFile,segmentations,exclude_overlap=False)
        print("Clustering Speakers")
        hardC, softC, centroids = self.pipeline.clustering(embeddings = embeddings,segmentations = segmentations)
        count = self.pipeline.speaker_count(
                    segmentations,
                    self.pipeline._segmentation.model.receptive_field,
                    warm_up=(0.0, 0.0),)
        print("Classifying Embeddings")
        embeddingClasses = np.zeros((embeddings.shape[0],embeddings.shape[1]))
        # Dumb loop to apply reclassifier
        # Counter for [silence,group] modifications for debugging
        tempCount = [0,0]
        for i,e in enumerate(embeddings):
            # Speaker, skip empty
            for j,eS in enumerate(e):
                if np.any(segmentations.data[i,:,j] > 0):
                    # Classify timestep
                    groupClass = self.classifyEmbedding(eS)
                    embeddingClasses[i][j] = groupClass
                    # Remove group from segmentations for later replacement
                    if groupClass == 2:
                        segmentations.data[i,:,j] = 0
                        tempCount[1] += 1
                    # Remove silence from segmentations if majority of timestep
                    elif groupClass == 0 and np.mean(segmentations.data[i,:,j]) < 0.5:
                        segmentations.data[i,:,j] = 0
                        tempCount[0] += 1
        print("Generating Annotation")
        #   shape: (num_chunks, num_speakers)
        # keep track of inactive speakers
        inactive_speakers = np.sum(segmentations.data, axis=1) == 0
        hardC[inactive_speakers] = -2
        discrete_diarization = self.pipeline.reconstruct(
            segmentations,
            hardC,
            count,)
        diarization = self.pipeline.to_annotation(
            discrete_diarization,
            min_duration_on=0.0,
            min_duration_off=self.pipeline.segmentation.min_duration_off,
        )
        # keep track of group speakers
        group_speakers = np.any(embeddingClasses >= 2,axis=1)
        # Seperate group speakers into 'pyannote.core.Segment's and apply to diarization under name 'group'
        start = None
        for timeStep in range(group_speakers.shape[0]):
            if group_speakers[timeStep] > 0:
                if start is None:
                    start = timeStep
            elif start is not None:
                segment = Segment(start, timeStep)
                diarization[segment] = 'group'
                start = None
        # Catch end case
        if start is not None:
            segment = Segment(start, embeddingClasses.shape[0]-1)
            diarization[segment] = 'group'

        # Estimate length of audio for charting
        totalTimeInSeconds = int(waveform.shape[-1]/sample_rate)

        # Rename labels to standard format
        currId = 0
        mapping = {}
        for label in diarization.labels():
            if label == 'group':
                continue
            else:
                currId += 1
                newLabel = f'SPEAKER_{currId:03d}'
                mapping[label] = newLabel
        diarization = diarization.rename_labels(mapping)
        print("Time in seconds calculated")
        return diarization, totalTimeInSeconds, waveform, sample_rate


    def activeSpeaker(self,inAnnotation,step=1):
        '''
        Determines the single active speaker for each timestep

        This estimates the primary speaker for each timestep for later use in identifying group discussion
        and when presenters/instructors change

        Parameters
        ----------
        inAnnotation : pyannote.core.Annotation
            Annotation object (diarization) to analyze
        step : float or int
            Time in seconds to use for determining current active speaker
        
        Returns
        -------
        speakerAtStep : list
            List of active speaker at each timestep (length matches stepTime)
        stepTime : list
            List of start time for each timestep (length matches speakerAtStep)
        speakerHierarchy : list
            List of speakers in order of priority. From most speech to least speech with group speech in front
        '''
        speakerAtStep = [None]
        stepTime = [0]
        speakerHierarchy = [label for label,_ in inAnnotation.chart()]

        # Identify hierarchy of speakers based on time present, with group in front
        for label in speakerHierarchy:
            # Move group labels to beginning of hierarchy (99 is training group code)
            if label == 'group' or label == 99:
                speakerHierarchy.remove(label)
                speakerHierarchy.insert(0,label)
        # Iterate over segments
        for segment,_,label in inAnnotation.itertracks(yield_label=True):
            startI = int(segment.start / step)
            # Lazy end assumption, always assumes one more step
            endI = int(segment.end / step) + 1
            # If stepTime and speakerAtStep not long enough for segment, then expand them
            while len(stepTime) < endI+1:
                stepTime.append(stepTime[-1]+step)
                speakerAtStep.append(None)
            # For each timestep in current segment, check current speaker against previous speaker
            for i in range(startI,endI+1):
                # No active speaker yet, so apply self
                if speakerAtStep[i] is None:
                    speakerAtStep[i] = label
                # If active speaker exists, check against hierarchy
                else:
                    currHier = speakerHierarchy.index(speakerAtStep[i])
                    newHier = speakerHierarchy.index(label)
                    if newHier < currHier:
                        speakerAtStep[i] = label
        return speakerAtStep, stepTime, speakerHierarchy
    
    def annotationToNoiseList(self,inAnnotation,maxTime,stepSize=2,windowSize=90):
        '''
        Determines which noise category applies to each timestep

        ...

        Parameters
        ----------
        inAnnotation : pyannote.core.Annotation
            Annotation object (diarization) to analyze
        maxTime : float
            Time at end of audio
        stepSize : float or int
            Time in seconds to use for determining current active speaker
        windowSize : float or int
            Time in seconds to use as window for determining noise categories

        Returns
        -------
        categorySegmentList : List[3,:]
            List of 3 Lists representing categories: group, individual, silence.  Each sublist contains tuples
            of (members,pyannote.core.Segment) where members is a string representing speaker names of all 
            relevant to the given Segment. Groups can contain '+' as a delimeter between speakers, e.g., 
            speaker1+speaker2+speaker3. Members are always in alphanumerical order.
        st : list
            List of start time for each timestep
        '''
        # Determine the active speaker for each timestep
        sas, st, sh = self.activeSpeaker(inAnnotation,step=stepSize)

        # Number of steps in window
        windowStepCount = windowSize / stepSize
        # Aggregate of scores for a given window
        timeStepAggregate = []
        # Class for a given window
        timeStepClass = []
        # Members for a given window
        timeStepMembers = []
        categories = ['group','individual','silence']
        # Initialize scores
        for i in st:
            timeStepAggregate.append({'individual':0,'group':0,'silence':0})
        # For each timestep
        for i,_ in enumerate(sas):
            decision = None
            groupCount = 0
            individuals = set()
            silenceCount = 0
            end = min(i+windowSize,len(sas))
            memberSet = set()
            # Iterate over timestep
            for j in range(i,end):
                # Add current speaker to member set
                if sas[j] is not None:
                    memberSet.add(sas[i])
                # Increase silence score if no speaker
                if sas[j] is None:
                    silenceCount += 1
                # Increase group score for known group IDs
                elif sas[j] == 'group' or sas[j] == 99:
                    groupCount += 1
                # TODO: Could probably replace individuals with memberSet, but leaving this code in for now
                else:
                    individuals.add(sas[j])
            # If majority of window is silence, then classify as silence
            if silenceCount > windowStepCount / 2:
                decision = 'silence'
            # If majority of window is known groups OR total individual speakers above threshold, then group
            elif sas[i] == 'group' or groupCount > windowStepCount / 2 or len(individuals) > 2:
                decision = 'group'
            # Classify as individual if not silence or group
            else:
                decision = 'individual'
                # If treated as individual, group should NOT be included!
                if 'group' in memberSet:
                    memberSet.remove('group')
            # Apply decision as score to all in window
            for j in range(i,end):
                timeStepAggregate[j][decision] += 1
            # Convert to list and sort for convenience
            memberSet = list(memberSet)
            memberSet.sort()
            timeStepMembers.append(memberSet)

        # Iterate over aggregate scores for each window
        for i,item in enumerate(timeStepAggregate):
            cat = None
            # Shortcut areas known to be group or silence, as these are more definite
            if sas[i] == 'group':
                cat = 'group'
            elif sas[i] is None:
                cat = 'silence'
            else:
                # Final classification is highest aggregate score
                cat = max(item, key=item.get)
            timeStepClass.append(cat)
            
            # For group decisions, members include all in window
            if cat == 'group': 
                timeStepMembers[i] = '+'.join(timeStepMembers[i])
            # Assume current speaker is the "individual" voice during window
            elif cat == 'individual':
                # Individual is first one seen
                firstIndividual = sas[i]
                firstIndividualIndex = i
                # If none currently (at/near end of silence) then find first individual instead
                while firstIndividual is None or firstIndividual == 'group':
                    firstIndividual = firstIndividual[firstIndividualIndex+1]
                    firstIndividualIndex += 1
                timeStepMembers[i] = firstIndividual
            # Remove all members if silence
            else:
                timeStepMembers[i] = None
                
        # For debug purposes
        singleDimList = []
        
        endTime = 0
        # [group list, individual list, silence list]
        categorySegmentList = []
        for c in categories:
            currList = []
            start = None
            currMembers = None
            duration = 0
            tracking = False
            # Iterate over timestep to generate pyannote.core.Segment foreach classification and member(s)
            for stepClass,timeIncrement,members in zip(timeStepClass,st,timeStepMembers):
                # Check for case of exact end of audio
                if st == maxTime:
                    continue
                # If current timestep classifies for current categorySegment list
                if stepClass == c:
                    # If we see the exact same member(s), then increment time
                    if currMembers == members:
                        duration += min(stepSize,maxTime-timeIncrement)
                    else:
                        # If already tracking, then generate Segment and restart tracking
                        if tracking:
                            singleDimList.append((currMembers,Segment(start,start+duration)))
                            currList.append((currMembers,Segment(start,start+duration)))
                            if start+duration > endTime:
                                endTime = start+duration
                            start = None
                            currMembers is None
                            duration = 0
                            tracking = False
                        start = timeIncrement
                        currMembers = members
                        duration += min(stepSize,maxTime-timeIncrement)
                        tracking = True
                # Timestep does NOT belong to current categorySegment list
                else:
                    # If tracking, then generate Segment and stop tracking
                    if tracking:
                        singleDimList.append((currMembers,Segment(start,start+duration)))
                        currList.append((currMembers,Segment(start,start+duration)))
                        if start+duration > endTime:
                            endTime = start+duration
                        start = None
                        currMembers is None
                        duration = 0
                        tracking = False
            # Exit case, if still tracking then generate final Segment
            if tracking:
                singleDimList.append((currMembers,Segment(start,start+duration)))
                currList.append((currMembers,Segment(start,start+duration)))
                if start+duration > endTime:
                    endTime = start+duration
            categorySegmentList.append(currList)
        # If we didn't end exactly on time, then fill in remaining time with Silence
        if endTime != maxTime:
            singleDimList.append((None,Segment(endTime,maxTime)))
            categorySegmentList[2].append((None,Segment(endTime,maxTime)))

        # For debug
        singleDimList = sorted(singleDimList,key=lambda index : index[1].start)
        print(singleDimList)
        return categorySegmentList, st
    
    def __call__(self,audioPath):
        '''
        Apply Sonogram to a given audio file

        ...
        
        Parameters
        ----------
        audioPath : str
            Full path to audio file to process

        Returns
        -------
        Returns
        -------
        annotation : pyannote.core.Annotation
            Diarization result from pipeline
        totalTimeInSeconds : int
            Approximate length of audio for use in diagrams
        waveform : np.array
            Audio waveform as loaded from file
        sampleRate : int
            Sample rate of loaded audio
        '''
        annotation, totalTimeInSeconds, waveform, sampleRate = self.processFile(audioPath)

        return annotation, totalTimeInSeconds, waveform, sampleRate

    def toDevice(self):
        '''
        Move Sonogram pipeline to device for accelerated processing

        ...
        '''
        self.pipeline.to(self.device)
        print(f"Sonogram moved to {self.device}")

    def toCPU(self):
        '''
        Move Sonogram pipeline to CPU to free up device space
        
        ...
        '''
        self.pipeline.to(torch.device('cpu'))
        print(f"Sonogram moved to CPU")