Update modules/utils/subtitle_manager.py

modules/utils/subtitle_manager.py

@@ -80,22 +80,64 @@ def get_vtt(segments):
     return output
 
 
+#def get_txt(segments):
+#    output = ""
+#    for i, segment in enumerate(segments):
+#        if segment['text'].startswith(' '):
+#            segment['text'] = segment['text'][1:]
+#        #output += f"{segment['text']}\n"
+#        output += f"{timeformat_txt(segment['start'])}\t{segment['text']}\n"
+#    return output
+
+#def get_plaintext(segments):
+#    output = ""
+#    for i, segment in enumerate(segments):
+#        if segment['text'].startswith(' '):
+#            segment['text'] = segment['text'][1:]
+#        output += f"{segment['text']}\n"
+#    return output
+
 def get_txt(segments):
+    bDiarization = False 
     output = ""
+
+    # Check if speakers are identified
     for i, segment in enumerate(segments):
-        if segment['text'].startswith(' '):
-            segment['text'] = segment['text'][1:]
-        #output += f"{segment['text']}\n"
-        output += f"{timeformat_txt(segment['start'])}\t{segment['text']}\n"
-    return output
+        if re.search(r'SPEAKER [0-9][0-9]: ',segment['text']) != None:
+            bDiarization = True
+            break
+
+    for i, segment in enumerate(segments):
+        if bDiarization:
+            speaker_id = ((segment['text']).split(":", 1)[0]).strip()
+            speaker_text = ((segment['text']).split(":", 1)[1]).strip()
+            output += f"{timeformat_txt(segment['start'])}\t{speaker_id}\t{speaker_text}\n"
+        else:
+            speaker_text = (segment['text']).strip()
+            output += f"{timeformat_txt(segment['start'])}\t{speaker_text}\n"
+
+    return output.rstrip("\n")
 
 def get_plaintext(segments):
+    bDiarization = False 
     output = ""
+
+    # Check if speakers are identified
     for i, segment in enumerate(segments):
-        if segment['text'].startswith(' '):
-            segment['text'] = segment['text'][1:]
-        output += f"{segment['text']}\n"
-    return output
+        if re.search(r'SPEAKER [0-9][0-9]: ',segment['text']) != None:
+            bDiarization = True
+            break
+
+    for i, segment in enumerate(segments):
+        if bDiarization:
+            speaker_id = ((segment['text']).split(":", 1)[0]).strip()
+            speaker_text = ((segment['text']).split(":", 1)[1]).strip()
+            output += f"{timeformat_txt({speaker_id}\t{speaker_text}\n"
+        else:
+            speaker_text = (segment['text']).strip()
+            output += f"{speaker_text}\n"
+
+    return output.rstrip("\n")
 
 def parse_srt(file_path):
     """Reads SRT file and returns as dict"""