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(): def __init__(self,version='1.0'): self.earlyCleanup = True self.isTPU = False self.isGPU = False try: 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") self.device = torch.device("cuda" if self.isGPU else "cpu") print(f"Using {self.device} instead.") self.version = version if version == 'speaker-diarization-3.1': self.pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization-3.1") 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 # Should manage this outside class now #self.pipeline.to(self.device) # Load SVM classifier 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)).item()) def processFile(self,filePath): # 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) # Rename labels 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): 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,maxTime,stepSize=2,windowSize=90): sas, st, sh = self.activeSpeaker(inAnnotation,step=stepSize) timeStepAggregate = [] timeStepClass = [] timeStepMembers = [] 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)) memberSet = set() for j in range(i,end): if sas[j] is not None: memberSet.add(sas[i]) 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) > 2: decision = 'group' else: decision = 'individual' 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) for i,item in enumerate(timeStepAggregate): 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]) elif cat == 'individual': timeStepMembers[i] = timeStepMembers[i][0] else: timeStepMembers[i] = None # For debug endTime = 0 singleDimList = [] categorySegmentList = [] for c in categories: currList = [] start = None currMembers = None duration = 0 tracking = False for stepClass,timeIncrement,members in zip(timeStepClass,st,timeStepMembers): # Check for case of exact end of audio if st == maxTime: continue if stepClass == c: if currMembers == members: duration += min(stepSize,maxTime-timeIncrement) else: 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 == None duration = 0 tracking = False start = timeIncrement currMembers = members duration += min(stepSize,maxTime-timeIncrement) tracking = True else: 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 == None duration = 0 tracking = False 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 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): annotation, totalTimeInSeconds, waveform, sampleRate = self.processFile(audioPath) return annotation, totalTimeInSeconds, waveform, sampleRate def toDevice(self): self.pipeline.to(self.device) print(f"Sonogram moved to {self.device}") def toCPU(self): self.pipeline.to(torch.device('cpu')) print(f"Sonogram moved to CPU")