Sonogram class overhaul

Updated Sonogram class to use trained model with reclassifier SVM and provide categorical classification

sonogram.py

@@ -2,22 +2,13 @@ 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
 
 class Sonogram():
 
-    def __init__(self,enableDenoise=False):
-        '''
-        Initialize Sonogram Class
-
-        enableDenoise : False|True
-            Legacy code to support denoise, which has currently been removed.  Consider removing if denoise will not be reimplemented in the future.
-        '''
-        #TODO: Should these be adjustable via initialization, or constants?
-        self.secondDifference = 5
-        self.gainWindow = 4
-        self.minimumGain = -45
-        self.maximumGain = -5
-        self.attenLimDB = 3
+    def __init__(self):
         self.earlyCleanup = True
         
         self.isTPU = False
@@ -40,86 +31,167 @@ class Sonogram():
         self.pipeline.to(self.device)
 
         # Load SVM classifier
-        with open('groupClassifier.pkl', 'rb') as f:
+        with open('05062026_groupClassifier.pkl', 'rb') as f:
             self.groupClassifier = pickle.load(f)
+
+    def classifyEmbedding(self,embedding):
+        return int(self.groupClassifier.predict(embedding.reshape(1, -1)))
     
     def processFile(self,filePath):
-        '''
-        Processes audio file to generate diarization output
-
-        filePath : string
-            Path to the audio file
-
-        Returns
-        --------
-        diarizationOutput : DiarizeOutput 
-            found here https://github.com/pyannote/pyannote-audio/blob/main/src/pyannote/audio/pipelines/speaker_diarization.py#L64
-
-        totalTimeInSeconds : int
-            Approximate total seconds of audio file
-
-        waveformGainAdjusted : np.array
-            The waveform of the audio file after equalization
-
-        sampleRate : int
-            The sample rate of the audio file
-        '''
-        print(f"Loading file : {filePath}")
-        waveformList, sampleRate = su.splitIntoTimeSegments(filePath,600)
-        print("File loaded")
-        waveformEnhanced = su.combineWaveforms(waveformList)    
-        if (self.earlyCleanup):
-            del waveformList
-        print("Equalizing Audio")
-        waveform_gain_adjusted = su.equalizeVolume()(waveformEnhanced,sampleRate,self.gainWindow,self.minimumGain,self.maximumGain)
-        if (self.earlyCleanup):
-            del waveformEnhanced
-        print("Audio Equalized")
-        print("Detecting speakers")
-        diarizationOutput, embeddings = self.pipeline({"waveform": waveform_gain_adjusted, "sample_rate": sampleRate}, return_embeddings=True)
-        annotations = diarizationOutput.speaker_diarization
-        embeddings = diarizationOutput.speaker_embeddings
-        print("Speakers Detected")
-        totalTimeInSeconds = int(waveform_gain_adjusted.shape[-1]/sampleRate)
+        # 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}
+        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
+        # Timestep
+        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):
+                    groupClass = self.classifyEmbedding(eS)
+                    embeddingClasses[i][j] = groupClass
+                    # Correct silence detections, remove group for later addition
+                    if groupClass == 2:
+                        segmentations.data[i,:,j] = 0
+                        tempCount[1] += 1
+                    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
+        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'
+            
+        totalTimeInSeconds = int(waveform.shape[-1]/sample_rate)
         print("Time in seconds calculated")
-        return diarizationOutput, totalTimeInSeconds, waveformGainAdjusted, sampleRate
-
-    def __call__(self,audioPath):
-        '''
-        Processes audio file to generate results necessary for app
+        return diarization, totalTimeInSeconds, waveform, sample_rate
 
-        filePath : string
-            Path to the audio file
 
-        Returns
-        --------
-        annotation : pyannote.core.annotation 
-            found here https://pyannote.github.io/pyannote-core/_modules/pyannote/core/annotation.html
-
-        totalTimeInSeconds : int
-            Approximate total seconds of audio file
-
-        waveformGainAdjusted : np.array
-            The waveform of the audio file after equalization
+    def activeSpeaker(self,inAnnotation,step=1):
+        speakerAtStep = [None]
+        stepTime = [0]
+        speakerHierarchy = [label for label,_ in inAnnotation.chart()]
+        
+        for label in speakerHierarchy:
+            # Move group labels to beginning of hierarchy
+            if label == 'group' or label == 99:
+                speakerHierarchy.remove(label)
+                speakerHierarchy.insert(0,label)
+        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
+            while len(stepTime) < endI+1:
+                stepTime.append(stepTime[-1]+step)
+                speakerAtStep.append(None)
+            for i in range(startI,endI+1):
+                if speakerAtStep[i] == None:
+                    speakerAtStep[i] = label
+                else:
+                    currHier = speakerHierarchy.index(speakerAtStep[i])
+                    newHier = speakerHierarchy.index(label)
+                    if newHier < currHier:
+                        speakerAtStep[i] = label
+        return speakerAtStep, stepTime, speakerHierarchy
+    
+    def annotationToNoiseList(self,inAnnotation,stepSize=2,windowSize=90):
+        sas, st, sh = self.activeSpeaker(inAnnotation,step=stepSize)
+        timeStepAggregate = []
+        timeStepClass = []
+        categories = ['group','individual','silence']
+        for i in st:
+            timeStepAggregate.append({'individual':0,'group':0,'silence':0})
+        for i,_ in enumerate(sas):
+            decision = None
+            groupCount = 0
+            individuals = set()
+            silenceCount = 0
+            end = min(i+windowSize,len(sas))
+            for j in range(i,end):
+                if sas[j] == None:
+                    silenceCount += 1
+                elif sas[j] == 'group' or sas[j] == 99:
+                    groupCount += 1
+                else:
+                    individuals.add(sas[j])
+            if silenceCount > windowSize / 2:
+                decision = 'silence'
+            elif groupCount > windowSize / 2 or len(individuals) > 3:
+                decision = 'group'
+            else:
+                decision = 'individual'
+            for j in range(i,end):
+                timeStepAggregate[j][decision] += 1
+        for item in timeStepAggregate:
+            timeStepClass.append(max(item, key=item.get))
 
-        sampleRate : int
-            The sample rate of the audio file
-        '''
-        diarizationOutput, totalTimeInSeconds, waveformGainAdjusted, sampleRate = self.processFile(audioPath)
-        annotation = diarizationOutput.speaker_diarization
+        
+        categorySegmentList = []
+        for c in categories:
+            currList = []
+            start = None
+            duration = 0
+            tracking = False
+            for stepClass,timeIncrement in zip(timeStepClass,st):
+                if stepClass == c:
+                    if start == None:
+                        start = timeIncrement
+                        duration += stepSize
+                        tracking = True
+                    else:
+                        duration += stepSize
+                else:
+                    if tracking:
+                        currList.append((start,duration))
+                        start = None
+                        duration = 0
+                        tracking = False
+            if tracking:
+                currList.append((start,duration))
+            categorySegmentList.append(currList)
+        return categorySegmentList, st
+    
+    def __call__(self,audioPath):
+        annotation, totalTimeInSeconds, waveform, sampleRate = self.processFile(audioPath)
 
-        # Relabel any existing silence and group speakers
-        labelMapping = {}
-        for s, speaker in enumerate(output.speaker_diarization.labels()):
-            diarizationOutput.speaker_embeddings[s]
-            prediction = self.groupClassifier.predict(diarizationOutput.speaker_embeddings[s].reshape(1,-1))
-            if prediction == 0:
-                labelMapping[speaker] = "silence"
-            elif prediction == 2:
-                labelMapping[speaker] = "group"
-            else:
-                # May not be necessary, consider using to reformat default names away from SPEAKER_XX
-                labelMapping[speaker] = speaker
-        # Rename in place
-        annotation.rename_labels(labelMapping)
-        return annotation, totalTimeInSeconds, waveformGainAdjusted, sampleRate
+        return annotation, totalTimeInSeconds, waveform, sampleRate