From c9e698642f22e444c477b1ed0159934f75bb0d94 Mon Sep 17 00:00:00 2001 From: su Date: Thu, 7 Mar 2024 18:59:40 -0500 Subject: [PATCH] add speaker labels to subtitles #117 --- src/subsai/models/whisperX_model.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/subsai/models/whisperX_model.py b/src/subsai/models/whisperX_model.py index 232a4ed..4ee2b10 100644 --- a/src/subsai/models/whisperX_model.py +++ b/src/subsai/models/whisperX_model.py @@ -149,7 +149,7 @@ def transcribe(self, media_file) -> str: try: event = SSAEvent(start=pysubs2.make_time(s=word["start"]), end=pysubs2.make_time(s=word["end"]), name=segment["speaker"] if self.speaker_labels else "") - event.plaintext = word["word"].strip() + event.plaintext = segment["speaker"] + ": " + word["word"].strip() if self.speaker_labels else word["word"].strip() subs.append(event) except Exception as e: logging.warning(f"Something wrong with {word}") @@ -159,7 +159,7 @@ def transcribe(self, media_file) -> str: for segment in result['segments']: event = SSAEvent(start=pysubs2.make_time(s=segment["start"]), end=pysubs2.make_time(s=segment["end"]), name=segment["speaker"] if self.speaker_labels else "") - event.plaintext = segment["text"].strip() + event.plaintext = segment["speaker"] + ": "+ segment["text"].strip() if self.speaker_labels else segment["text"].strip() subs.append(event) else: raise Exception(f'Unknown `segment_type` value, it should be one of the following: '