Update voice_analysis.py

voice_analysis.py

@@ -22,9 +22,20 @@ def diarize_speakers(audio_path):
     
     pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization-3.1", use_auth_token=hf_token)
     diarization = pipeline(audio_path)
-    return diarization
     
-def get_speaker_embeddings(audio_path, diarization, model_name="pyannote/embedding"):
+    # Identify the most frequent speaker
+    speaker_segments = {}
+    for turn, _, speaker in diarization.itertracks(yield_label=True):
+        if speaker not in speaker_segments:
+            speaker_segments[speaker] = 0
+        speaker_segments[speaker] += turn.end - turn.start
+    
+    most_frequent_speaker = max(speaker_segments, key=speaker_segments.get)
+    
+    return diarization, most_frequent_speaker
+
+    
+def get_speaker_embeddings(audio_path, diarization, most_frequent_speaker, model_name="pyannote/embedding"):
     model = Model.from_pretrained(model_name, use_auth_token=os.environ.get("py_annote_hf_token"))
     waveform, sample_rate = torchaudio.load(audio_path)
     duration = waveform.shape[1] / sample_rate
@@ -39,6 +50,9 @@ def get_speaker_embeddings(audio_path, diarization, model_name="pyannote/embeddi
     
     embeddings = []
     for turn, _, speaker in diarization.itertracks(yield_label=True):
+        if speaker != most_frequent_speaker:
+            continue
+        
         start_frame = int(turn.start * sample_rate)
         end_frame = int(turn.end * sample_rate)
         segment = waveform[:, start_frame:end_frame]
@@ -79,6 +93,18 @@ def get_speaker_embeddings(audio_path, diarization, model_name="pyannote/embeddi
     
     return embeddings, duration
 
+    
+    # Ensure embeddings cover the entire duration
+    if embeddings and embeddings[-1]['time'] + embeddings[-1]['duration'] < duration:
+        embeddings.append({
+            "time": duration,
+            "duration": 0,
+            "embedding": np.zeros_like(embeddings[0]['embedding']),
+            "speaker": "silence"
+        })
+    
+    return embeddings, duration
+
 def align_voice_embeddings(voice_embeddings, frame_count, fps, audio_duration):
     aligned_embeddings = []
     current_embedding_index = 0