add speaker labels to subtitles #117

absadiki · Mar 7, 2024 · c9e6986 · c9e6986
1 parent 38959f8
commit c9e6986
Showing 1 changed file with 2 additions and 2 deletions.
diff --git 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: '