diff --git a/whisper_live/client.py b/whisper_live/client.py index 96971c9..d61e4b5 100644 --- a/whisper_live/client.py +++ b/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): - for seg in message: + n_segments = len(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: """