From d18bcd4c210b24396f38bffb0b46718e62806abf Mon Sep 17 00:00:00 2001 From: jhj0517 <97279763+jhj0517@users.noreply.github.com> Date: Sat, 2 Nov 2024 14:56:07 +0900 Subject: [PATCH] Fix diarization type --- modules/diarize/diarize_pipeline.py | 6 +++--- modules/diarize/diarizer.py | 1 - 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/modules/diarize/diarize_pipeline.py b/modules/diarize/diarize_pipeline.py index bdb7d53..4313f5c 100644 --- 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 index 4c0727a..38e150a 100644 --- 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"]