Merge pull request #374 from jhj0517/fix/diarization

Fix diarization
jhj0517 · Nov 2, 2024 · f12a40c · f12a40c
2 parents 0295d69 + d18bcd4
commit f12a40c
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 4 deletions.
diff --git a/modules/diarize/diarize_pipeline.py b/modules/diarize/diarize_pipeline.py
@@ -44,9 +44,9 @@ def __call__(self, audio: Union[str, np.ndarray], min_speakers=None, max_speaker
 
 def assign_word_speakers(diarize_df, transcript_result, fill_nearest=False):
     transcript_segments = transcript_result["segments"]
+    if transcript_segments and isinstance(transcript_segments[0], Segment):
+        transcript_segments = [seg.model_dump() for seg in transcript_segments]
     for seg in transcript_segments:
-        if isinstance(seg, Segment):
-            seg = seg.model_dump()
         # assign speaker to segment (if any)
         diarize_df['intersection'] = np.minimum(diarize_df['end'], seg['end']) - np.maximum(diarize_df['start'],
                                                                                             seg['start'])
@@ -88,7 +88,7 @@ def assign_word_speakers(diarize_df, transcript_result, fill_nearest=False):
                     if word_speaker is not None:
                         word["speaker"] = word_speaker
 
-    return transcript_result
+    return {"segments": transcript_segments}
 
 
 class DiarizationSegment:

diff --git a/modules/diarize/diarizer.py b/modules/diarize/diarizer.py
@@ -71,7 +71,6 @@ def run(self,
 
         segments_result = []
         for segment in diarized_result["segments"]:
-            segment = segment.dict()
             speaker = "None"
             if "speaker" in segment:
                 speaker = segment["speaker"]