Update modules/utils/subtitle_manager.py

modules/utils/subtitle_manager.py

@@ -32,16 +32,30 @@ def get_srt(segments):
         output += f"{segment['text']}\n\n"
     return output
 
-def get_csv(segments,diarization: bool = False):
-    if diarization:
-        output = "Line;Start time;End time;Speaker;Text;\n"
-    else:
-        output = "Line;Start time;End time;Text;\n"
-        
+def get_csv(segments):
+    bDiarization = False    
+    output = ""
+            
     for i, segment in enumerate(segments):
         if segment['text'].startswith(' '):
-            segment['text'] = segment['text'][1:]            
-        output += f"{i + 1};{timeformat_srt(segment['start'])};{timeformat_srt(segment['end'])};{segment['text']};\n"
+            segment['text'] = segment['text'][1:]
+
+        # Check if speakers are identified and get speaker id & text
+        temp = re.search(r'SPEAKER_[0-9][0-9]: ',segment['text'])
+        if temp != None:
+        	speaker_id = (temp.group()).replace(': ','')
+        	speaker_text = test_input.replace(temp.group(),'')
+            output += f"{i + 1};{timeformat_srt(segment['start'])};{timeformat_srt(segment['end'])};{speaker_id};{speaker_text};\n"
+            bDiarization = True
+        else:
+            output += f"{i + 1};{timeformat_srt(segment['start'])};{timeformat_srt(segment['end'])};{segment['text']};\n"
+
+        # Add titles to csv file
+        if bDiarization:
+            output = "Line;Start time;End time;Speaker;Text;\n" + output
+        else:    
+            output = "Line;Start time;End time;Text;\n" + output
+
     return output
 
 def get_vtt(segments):