add: save_transcript to srt file

whisper_live/client.py

@@ -13,6 +13,29 @@ import uuid
 import time
 
 
+def format_time(s):
+    """Convert seconds (float) to SRT time format."""
+    hours = int(s // 3600)
+    minutes = int((s % 3600) // 60)
+    seconds = int(s % 60)
+    milliseconds = int((s - int(s)) * 1000)
+    return f"{hours:02}:{minutes:02}:{seconds:02},{milliseconds:03}"
+
+def create_srt_file(segments, output_file):
+    with open(output_file, 'w', encoding='utf-8') as srt_file:
+        segment_number = 1
+        for segment in segments:
+            start_time = format_time(float(segment['start']))
+            end_time = format_time(float(segment['end']))
+            text = segment['text']
+            
+            srt_file.write(f"{segment_number}\n")
+            srt_file.write(f"{start_time} --> {end_time}\n")
+            srt_file.write(f"{text}\n\n")
+            
+            segment_number += 1
+
+
 def resample(file: str, sr: int = 16000):
     """
     # https://github.com/openai/whisper/blob/7858aa9c08d98f75575035ecd6481f462d66ca27/whisper/audio.py#L22
@@ -57,6 +80,7 @@ class Client:
         lang=None,
         translate=False,
         model="small",
+        srt_file_path="output.srt"
     ):
         """
         Initializes a Client instance for audio recording and streaming to a server.
@@ -89,6 +113,7 @@ class Client:
         self.language = lang
         self.model = model
         self.server_error = False
+        self.srt_file_path = srt_file_path
 
         if translate:
             self.task = "translate"
@@ -127,6 +152,7 @@ class Client:
         self.ws_thread.start()
 
         self.frames = b""
+        self.transcript = []
         print("[INFO]: * recording")
 
     def on_message(self, ws, message):
@@ -181,12 +207,21 @@ class Client:
 
         message = message["segments"]
         text = []
-        if len(message):
+        n_segments = len(message)
-            for seg in message:
+
+        if n_segments:
+            for i, seg in enumerate(message):
                 if text and text[-1] == seg["text"]:
                     # already got it
                     continue
                 text.append(seg["text"])
+
+                if i == n_segments-1: 
+                    self.last_segment = seg
+                else:
+                    if not len(self.transcript) or float(seg['start']) >= float(self.transcript[-1]['end']):
+                        self.transcript.append(seg)
+
         # keep only last 3
         if len(text) > 3:
             text = text[-3:]
@@ -302,6 +337,7 @@ class Client:
                 assert self.last_response_recieved
                 while time.time() - self.last_response_recieved < self.disconnect_if_no_response_for:
                     continue
+                self.write_srt_file(self.srt_file_path)
                 self.stream.close()
                 self.close_websocket()
 
@@ -311,6 +347,7 @@ class Client:
                 self.stream.close()
                 self.p.terminate()
                 self.close_websocket()
+                self.write_srt_file(self.srt_file_path)
                 print("[INFO]: Keyboard interrupt.")
 
     def close_websocket(self):
@@ -438,6 +475,7 @@ class Client:
                     t.start()
                     n_audio_file += 1
                     self.frames = b""
+            self.write_srt_file(self.srt_file_path)
 
         except KeyboardInterrupt:
             if len(self.frames):
@@ -451,6 +489,7 @@ class Client:
             self.close_websocket()
 
             self.write_output_recording(n_audio_file, out_file)
+            self.write_srt_file(self.srt_file_path)
 
     def write_output_recording(self, n_audio_file, out_file):
         """
@@ -487,6 +526,9 @@ class Client:
                 os.remove(in_file)
         wavfile.close()
 
+    def write_srt_file(self, output_path="output.srt"):
+        self.transcript.append(self.last_segment)
+        create_srt_file(self.transcript, output_path)
 
 class TranscriptionClient:
     """