From a5356509b6fc4465db6ceeacadada9c2bdcbd34b Mon Sep 17 00:00:00 2001 From: Marco Vela Date: Tue, 10 Oct 2023 14:50:41 -0500 Subject: [PATCH] fix(diarize): key error on empty track --- whisperx/diarize.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/whisperx/diarize.py b/whisperx/diarize.py index d59a847..55fe220 100644 --- a/whisperx/diarize.py +++ b/whisperx/diarize.py @@ -26,10 +26,9 @@ class DiarizationPipeline: 'sample_rate': SAMPLE_RATE } segments = self.model(audio_data, min_speakers=min_speakers, max_speakers=max_speakers) - diarize_df = pd.DataFrame(segments.itertracks(yield_label=True)) - diarize_df['start'] = diarize_df[0].apply(lambda x: x.start) - diarize_df['end'] = diarize_df[0].apply(lambda x: x.end) - diarize_df.rename(columns={2: "speaker"}, inplace=True) + diarize_df = pd.DataFrame(segments.itertracks(yield_label=True), columns=['segment', 'label', 'speaker']) + diarize_df['start'] = diarize_df['segment'].apply(lambda x: x.start) + diarize_df['end'] = diarize_df['segment'].apply(lambda x: x.end) return diarize_df