From a65d69cfbd328bbc4aaae553137541fdb412fbe7 Mon Sep 17 00:00:00 2001 From: Ryan Pavlik Date: Fri, 22 Sep 2023 14:30:52 -0500 Subject: [PATCH] Tidy up the client python file --- whisper_live/client.py | 289 ++++++++++++++++++++++------------------- 1 file changed, 158 insertions(+), 131 deletions(-) diff --git a/whisper_live/client.py b/whisper_live/client.py index 663f579..19e856a 100644 --- a/whisper_live/client.py +++ b/whisper_live/client.py @@ -1,5 +1,4 @@ import os -import argparse import wave import numpy as np @@ -44,47 +43,58 @@ def resample(file: str, sr: int = 16000): class Client: - CHUNK = 1024 - FORMAT = pyaudio.paInt16 - CHANNELS = 1 - RATE = 16000 - RECORD_SECONDS = 60000 - RECORDING = False - multilingual = False - language = None - task = "transcribe" - uid = str(uuid.uuid4()) - WAITING = False - LAST_RESPONSE_RECIEVED = None - DISCONNECT_IF_NO_RESPONSE_FOR = 15 - - def __init__(self, host=None, port=None, is_multilingual=False, lang=None, translate=False): - Client.multilingual = is_multilingual - Client.language = lang if is_multilingual else "en" + INSTANCES = {} + + def __init__( + self, host=None, port=None, is_multilingual=False, lang=None, translate=False + ): + self.chunk = 1024 + self.format = pyaudio.paInt16 + self.channels = 1 + self.rate = 16000 + self.record_seconds = 60000 + self.recording = False + self.multilingual = False + self.language = None + self.task = "transcribe" + self.uid = str(uuid.uuid4()) + self.waiting = False + self.last_response_recieved = None + self.disconnect_if_no_response_for = 15 + self.multilingual = is_multilingual + self.language = lang if is_multilingual else "en" if translate: - Client.task = "translate" + self.task = "translate" self.timestamp_offset = 0.0 self.audio_bytes = None self.p = pyaudio.PyAudio() - self.stream = self.p.open(format=self.FORMAT, - channels=self.CHANNELS, - rate=self.RATE, - input=True, - frames_per_buffer=self.CHUNK) + self.stream = self.p.open( + format=self.format, + channels=self.channels, + rate=self.rate, + input=True, + frames_per_buffer=self.chunk, + ) # create websocket connection if host is not None and port is not None: - socket_url = f"ws://{host}:{port}" - self.client_socket = websocket.WebSocketApp(socket_url, - on_open=Client.on_open, - on_message=Client.on_message, - on_error=Client.on_error, - on_close=Client.on_close) + socket_url = f"ws://{host}:{port}" + self.client_socket = websocket.WebSocketApp( + socket_url, + on_open=lambda ws: self.on_open(ws), + on_message=lambda ws, message: self.on_message(ws, message), + on_error=lambda ws, error: self.on_error(ws, error), + on_close=lambda ws, close_status_code, close_msg: self.on_close( + ws, close_status_code, close_msg + ), + ) else: print("[ERROR]: No host or port specified.") return + Client.INSTANCES[self.uid] = self + # start websocket client in a thread self.ws_thread = threading.Thread(target=self.client_socket.run_forever) self.ws_thread.setDaemon(True) @@ -92,122 +102,130 @@ class Client: self.frames = b"" print("[INFO]: * recording") - - @staticmethod - def on_message(ws, message): - Client.LAST_RESPONSE_RECIEVED = time.time() + + def on_message(self, ws, message): + self.last_response_recieved = time.time() message = json.loads(message) - if message.get('uid')!=Client.uid: + + if self.uid != message.get("uid"): print("[ERROR]: invalid client uid") return - - if "status" in message.keys() and message["status"] == "WAIT": - Client.WAITING = True - print(f"[INFO]:Server is full. Estimated wait time {round(message['message'])} minutes.") - + + if "status" in message.keys() and message["status"] == "WAIT": + self.waiting = True + print( + f"[INFO]:Server is full. Estimated wait time {round(message['message'])} minutes." + ) + if "message" in message.keys() and message["message"] == "DISCONNECT": print("[INFO]: Server overtime disconnected.") - Client.RECORDING = False + self.recording = False if "message" in message.keys() and message["message"] == "SERVER_READY": - Client.RECORDING = True + self.recording = True return if "language" in message.keys(): - Client.language = message.get("language") + self.language = message.get("language") lang_prob = message.get("language_prob") - print(f"[INFO]: Server detected language {Client.language} with probability {lang_prob}") + print( + f"[INFO]: Server detected language {self.language} with probability {lang_prob}" + ) return if "segments" not in message.keys(): return - + message = message["segments"] text = [] if len(message): for seg in message: - if len(text): - if text[-1] != seg["text"]: - text.append(seg["text"]) - else: - text.append(seg["text"]) + if text and text[-1] == seg["text"]: + # already got it + continue + text.append(seg["text"]) + # keep only last 3 if len(text) > 3: text = text[-3:] wrapper = textwrap.TextWrapper(width=60) word_list = wrapper.wrap(text="".join(text)) # Print each line. - if os.name=='nt': - os.system('cls') + if os.name == "nt": + os.system("cls") else: - os.system('clear') + os.system("clear") for element in word_list: print(element) - @staticmethod - def on_error(ws, error): + def on_error(self, ws, error): print(error) - @staticmethod - def on_close(ws, close_status_code, close_msg): - print(f"[INFO]: Websocket connection closed.") + def on_close(self, ws, close_status_code, close_msg): + print(f"[INFO]: Websocket connection closed: {close_status_code}: {close_msg}") - @staticmethod - def on_open(ws): - print(Client.multilingual, Client.language, Client.task) + def on_open(self, ws): + print(self.multilingual, self.language, self.task) print("[INFO]: Opened connection") - ws.send(json.dumps({ - 'uid': Client.uid, - 'multilingual': Client.multilingual, - 'language': Client.language, - 'task': Client.task - })) + ws.send( + json.dumps( + { + "uid": self.uid, + "multilingual": self.multilingual, + "language": self.language, + "task": self.task, + } + ) + ) @staticmethod def bytes_to_float_array(audio_bytes): - raw_data = np.frombuffer( - buffer=audio_bytes, dtype=np.int16 - ) + raw_data = np.frombuffer(buffer=audio_bytes, dtype=np.int16) return raw_data.astype(np.float32) / 32768.0 - + def send_packet_to_server(self, message): try: self.client_socket.send(message, websocket.ABNF.OPCODE_BINARY) except Exception as e: print(e) - + def play_file(self, filename): # read audio and create pyaudio stream - self.wf = wave.open(filename, 'rb') - self.stream = self.p.open(format=self.p.get_format_from_width(self.wf.getsampwidth()), - channels=self.wf.getnchannels(), - rate=self.wf.getframerate(), + with wave.open(filename, "rb") as wavfile: + self.stream = self.p.open( + format=self.p.get_format_from_width(wavfile.getsampwidth()), + channels=wavfile.getnchannels(), + rate=wavfile.getframerate(), input=True, output=True, - frames_per_buffer=self.CHUNK) - try: - while Client.RECORDING: - data = self.wf.readframes(self.CHUNK) - if data==b'': break + frames_per_buffer=self.chunk, + ) + try: + while self.recording: + data = wavfile.readframes(self.chunk) + if data == b"": + break - audio_array = Client.bytes_to_float_array(data) - self.send_packet_to_server(audio_array.tobytes()) - self.stream.write(data) + audio_array = self.bytes_to_float_array(data) + self.send_packet_to_server(audio_array.tobytes()) + self.stream.write(data) - self.wf.close() - elapsed_time = time.time() - self.LAST_RESPONSE_RECIEVED - while elapsed_time < self.DISCONNECT_IF_NO_RESPONSE_FOR: - continue - self.stream.close() - self.close_websocket() + wavfile.close() - except KeyboardInterrupt: - self.wf.close() - self.stream.stop_stream() - self.stream.close() - self.p.terminate() - self.close_websocket() - print("[INFO]: Keyboard interrupt.") + assert self.last_response_recieved + elapsed_time = time.time() - self.last_response_recieved + while elapsed_time < self.disconnect_if_no_response_for: + continue + self.stream.close() + self.close_websocket() + + except KeyboardInterrupt: + wavfile.close() + self.stream.stop_stream() + self.stream.close() + self.p.terminate() + self.close_websocket() + print("[INFO]: Keyboard interrupt.") def close_websocket(self): try: @@ -222,14 +240,14 @@ class Client: def get_client_socket(self): return self.client_socket - + def write_audio_frames_to_file(self, frames, file_name): - wf = wave.open(file_name, 'wb') - wf.setnchannels(self.CHANNELS) - wf.setsampwidth(2) - wf.setframerate(self.RATE) - wf.writeframes(frames) - wf.close() + with wave.open(file_name, "wb") as wavfile: + wavfile: wave.Wave_write + wavfile.setnchannels(self.channels) + wavfile.setsampwidth(2) + wavfile.setframerate(self.rate) + wavfile.writeframes(frames) def record(self, out_file="output_recording.wav"): n_audio_file = 0 @@ -237,20 +255,24 @@ class Client: if not os.path.exists("chunks"): os.makedirs("chunks", exist_ok=True) try: - for _ in range(0, int(self.RATE / self.CHUNK * self.RECORD_SECONDS)): - if not Client.RECORDING: break - data = self.stream.read(self.CHUNK) + for _ in range(0, int(self.rate / self.chunk * self.record_seconds)): + if not self.recording: + break + data = self.stream.read(self.chunk) self.frames += data audio_array = Client.bytes_to_float_array(data) - + self.send_packet_to_server(audio_array.tobytes()) # save frames if more than a minute - if len(self.frames) > 60*self.RATE: + if len(self.frames) > 60 * self.rate: t = threading.Thread( target=self.write_audio_frames_to_file, - args=(self.frames[:], f"chunks/{n_audio_file}.wav", ) + args=( + self.frames[:], + f"chunks/{n_audio_file}.wav", + ), ) t.start() n_audio_file += 1 @@ -259,7 +281,8 @@ class Client: except KeyboardInterrupt: if len(self.frames): self.write_audio_frames_to_file( - self.frames[:], f"chunks/{n_audio_file}.wav") + self.frames[:], f"chunks/{n_audio_file}.wav" + ) n_audio_file += 1 self.stream.stop_stream() self.stream.close() @@ -268,33 +291,38 @@ class Client: # combine all the audio files self.write_output_recording(n_audio_file, out_file) - + def write_output_recording(self, n_audio_file, out_file): - input_files = [f"chunks/{i}.wav" for i in range(n_audio_file) if os.path.exists(f"chunks/{i}.wav")] - wf = wave.open(out_file, 'wb') - wf.setnchannels(self.CHANNELS) - wf.setsampwidth(2) - wf.setframerate(self.RATE) - for in_file in input_files: - w = wave.open(in_file, 'rb') - while True: - data = w.readframes(self.CHUNK) - if data==b'': break - wf.writeframes(data) - w.close() - # remove this file - os.remove(in_file) - wf.close() + input_files = [ + f"chunks/{i}.wav" + for i in range(n_audio_file) + if os.path.exists(f"chunks/{i}.wav") + ] + with wave.open(out_file, "wb") as wavfile: + wavfile: wave.Wave_write + wavfile.setnchannels(self.channels) + wavfile.setsampwidth(2) + wavfile.setframerate(self.rate) + for in_file in input_files: + with wave.open(in_file, "rb") as wav_in: + while True: + data = wav_in.readframes(self.chunk) + if data == b"": + break + wavfile.writeframes(data) + # remove this file + os.remove(in_file) + wavfile.close() class TranscriptionClient: def __init__(self, host, port, is_multilingual=False, lang=None, translate=False): self.client = Client(host, port, is_multilingual, lang, translate) - + def __call__(self, audio=None): print("[INFO]: Waiting for server ready ...") - while not Client.RECORDING: - if Client.WAITING: + while not self.client.recording: + if self.client.waiting: self.client.close_websocket() return pass @@ -304,4 +332,3 @@ class TranscriptionClient: self.client.play_file(resampled_file) else: self.client.record() -