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")