Documentation provided for Sonogram.py

sonogram.py

@@ -11,29 +11,76 @@ 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}")
-            # Fallback to CPU or other devices if needed
             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
@@ -51,25 +98,60 @@ class Sonogram():
             
             baselinePipeline.segmentation = ParamDict(min_duration_off=0.0,)
             self.pipeline = baselinePipeline
-        
-
-        # Should manage this outside class now
-        #self.pipeline.to(self.device)
 
-        # Load SVM classifier
+        # 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")
@@ -82,19 +164,21 @@ class Sonogram():
                     warm_up=(0.0, 0.0),)
         print("Classifying Embeddings")
         embeddingClasses = np.zeros((embeddings.shape[0],embeddings.shape[1]))
-        # Dumb loop
-        # Timestep
+        # 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
-                    # Correct silence detections, remove group for later addition
+                    # 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
@@ -114,6 +198,7 @@ class Sonogram():
         )
         # 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:
@@ -127,10 +212,11 @@ class Sonogram():
         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
+        # Rename labels to standard format
         currId = 0
         mapping = {}
         for label in diarization.labels():
@@ -146,25 +232,53 @@ class Sonogram():
 
 
     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
+            # 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] == None:
                     speakerAtStep[i] = label
+                # If active speaker exists, check against hierarchy
                 else:
                     currHier = speakerHierarchy.index(speakerAtStep[i])
                     newHier = speakerHierarchy.index(label)
@@ -173,14 +287,48 @@ class Sonogram():
         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
@@ -188,42 +336,60 @@ class Sonogram():
             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] == 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 silenceCount > windowSize / 2:
+            # If majority of window is silence, then classify as silence
+            if silenceCount > windowStepCount / 2:
                 decision = 'silence'
-            elif sas[i] == 'group' or groupCount > windowSize / 2 or len(individuals) > 2:
+            # 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):
+            # 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':
                 timeStepMembers[i] = sas[i]
+            # Remove all members if silence
             else:
                 timeStepMembers[i] = None
                 
-        # For debug
-        endTime = 0
+        # For debug purposes
         singleDimList = []
+        
+        endTime = 0
+        # [group list, individual list, silence list]
         categorySegmentList = []
         for c in categories:
             currList = []
@@ -231,15 +397,18 @@ class Sonogram():
             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)))
@@ -253,7 +422,9 @@ class Sonogram():
                         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)))
@@ -263,13 +434,14 @@ class Sonogram():
                         currMembers == 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)
-        # Check where we left off
+        # 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)))
@@ -280,14 +452,47 @@ class Sonogram():
         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")