Update diarization.py

diarization.py

@@ -97,15 +97,22 @@ def transcribe_audio(audio_path, language):
 
 def create_combined_srt(transcription_chunks, diarization, output_path):
     speaker_segments = []
-    speaker_map = {}
-    current_speaker_num = 1
+    speaker_durations = defaultdict(float)
 
+    # First pass: calculate durations
     for segment, _, speaker in diarization.itertracks(yield_label=True):
-        if speaker not in speaker_map:
-            speaker_map[speaker] = f"Speaker {current_speaker_num}"
-            current_speaker_num += 1
-        speaker_segments.append((segment.start, segment.end, speaker_map[speaker]))
+        speaker_durations[speaker] += segment.end - segment.start
+        speaker_segments.append((segment.start, segment.end, speaker))
+
+    # Sort speakers by duration
+    sorted_speakers = sorted(speaker_durations.items(), key=lambda x: x[1], reverse=True)
+    
+    # Create mapping of original labels to Speaker 1, Speaker 2, etc.
+    speaker_map = {}
+    for i, (speaker, _) in enumerate(sorted_speakers, start=1):
+        speaker_map[speaker] = f"Speaker {i}"
 
+    # Write the SRT content
     with open(output_path, 'w', encoding='utf-8') as srt_file:
         for i, chunk in enumerate(transcription_chunks, 1):
             start_time, end_time = chunk["start"], chunk["end"]
@@ -115,7 +122,7 @@ def create_combined_srt(transcription_chunks, diarization, output_path):
             current_speaker = "Unknown"
             for seg_start, seg_end, speaker in speaker_segments:
                 if seg_start <= start_time < seg_end:
-                    current_speaker = speaker
+                    current_speaker = speaker_map[speaker]
                     break
 
             # Format timecodes as h:mm:ss (without leading zeros for hours)
@@ -125,6 +132,12 @@ def create_combined_srt(transcription_chunks, diarization, output_path):
             srt_file.write(f"{i}\n")
             srt_file.write(f"{{{current_speaker}}}\n time: ({start_str} --> {end_str})\n text: {text}\n\n")
 
+    # Add information about the two most frequent speakers
+    with open(output_path, 'a', encoding='utf-8') as srt_file:
+        for i, (speaker, duration) in enumerate(sorted_speakers[:2], start=1):
+            duration_str = format_timestamp(duration).split('.')[0].lstrip('0')
+            srt_file.write(f"Speaker {i} (originally {speaker}): total duration {duration_str}\n")
+
     # Calculate speaker durations
     speaker_durations = defaultdict(float)
     for seg_start, seg_end, speaker in speaker_segments: