diff --git a/whisper_live/__init__.py b/whisper_live/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/client.py b/whisper_live/client.py similarity index 54% rename from client.py rename to whisper_live/client.py index 44fe89c..158e364 100644 --- a/client.py +++ b/whisper_live/client.py @@ -12,204 +12,6 @@ import json import websocket -CHUNK = 1024 -FORMAT = pyaudio.paInt16 -CHANNELS = 1 -RATE = 16000 -RECORD_SECONDS = 60000 -START_RECORDING = False -multilingual = False -language = None - - - -def on_message(ws, message): - global START_RECORDING, language - message = json.loads(message) - if message == "SERVER_READY": - START_RECORDING = True - return - - if isinstance(message, dict): - language = message.get("language") - lang_prob = message.get("language_prob") - print(f"Server detected language {language} with probability {lang_prob}") - return - - 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 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') - else: - os.system('clear') - for element in word_list: - print(element) - -def on_error(ws, error): - print(error) - -def on_close(ws, close_status_code, close_msg): - print("### websocket connection closed ###") - -def on_open(ws): - global multilingual, language, task - print(multilingual, language, task) - - print("Opened connection") - ws.send(json.dumps({ - 'multilingual': multilingual[0], - 'language': language[0], - 'task': task - })) - - -class Client: - def __init__(self, host=None, port=None): - self.timestamp_offset = 0.0 - self.audio_bytes = None - self.p = pyaudio.PyAudio() - self.stream = self.p.open(format=FORMAT, - channels=CHANNELS, - rate=RATE, - input=True, - frames_per_buffer=CHUNK) - print(self.p.get_sample_size(FORMAT)) - - # 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=on_open, - on_message=on_message, - on_error=on_error, - on_close=on_close) - else: - print("No host or port specified.") - return - - # start websocket client in a thread - self.ws_thread = threading.Thread(target=self.client_socket.run_forever) - self.ws_thread.setDaemon(True) - self.ws_thread.start() - - self.frames = b"" - print("* recording") - - def send_packet_to_server(self, message): - try: - self.client_socket.send(message, websocket.ABNF.OPCODE_BINARY) - except Exception as e: - print(e) - - @staticmethod - def bytes_to_float_array(audio_bytes): - raw_data = np.frombuffer( - buffer=audio_bytes, dtype=np.int16 - ) - return raw_data.astype(np.float32) / 32768.0 - - 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(), - input=True, - output=True, - frames_per_buffer=CHUNK) - try: - while True: - data = self.wf.readframes(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) - - self.wf.close() - self.stream.close() - - except KeyboardInterrupt: - print("Keyboard interrupt.") - - - 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(CHANNELS) - wf.setsampwidth(2) - wf.setframerate(RATE) - wf.writeframes(frames) - wf.close() - - def record(self, out_file="output_recording.wav"): - n_audio_file = 0 - # create dir for saving audio chunks - if not os.path.exists("chunks"): - os.makedirs("chunks", exist_ok=True) - try: - for _ in range(0, int(RATE / CHUNK * RECORD_SECONDS)): - data = self.stream.read(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*RATE: - t = threading.Thread( - target=self.write_audio_frames_to_file, - args=(self.frames[:], f"chunks/{n_audio_file}.wav", ) - ) - t.start() - n_audio_file += 1 - self.frames = b"" - - except KeyboardInterrupt: - if len(self.frames): - self.write_audio_frames_to_file( - self.frames[:], f"chunks/{n_audio_file}.wav") - n_audio_file += 1 - self.stream.stop_stream() - self.stream.close() - self.p.terminate() - - - # 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(CHANNELS) - wf.setsampwidth(2) - wf.setframerate(RATE) - for in_file in input_files: - w = wave.open(in_file, 'rb') - while True: - data = w.readframes(CHUNK) - if data==b'': break - wf.writeframes(data) - w.close() - # remove this file - os.remove(in_file) - wf.close() - - def resample(file: str, sr: int = 16000): """ # https://github.com/openai/whisper/blob/7858aa9c08d98f75575035ecd6481f462d66ca27/whisper/audio.py#L22 @@ -239,30 +41,235 @@ def resample(file: str, sr: int = 16000): return resampled_file -if __name__=="__main__": - parser = argparse.ArgumentParser() - parser.add_argument('--audio', type=str, help='audio file to transcribe') - parser.add_argument('--host', default=None, type=str, help='websocket server address to connect to') - parser.add_argument('--port', default=None, type=str, help='websocket server port to connect to') - parser.add_argument('--multilingual', action="store_true", help='use multilingual model') - parser.add_argument('--language', default=None, type=str, help='languages to use') - parser.add_argument( - '--task', default="transcribe", type=str, help='task transcribe/translate (translates from any to english)') - opt = parser.parse_args() - print(opt) - multilingual=opt.multilingual, - language = opt.language if opt.multilingual else "en", - task = opt.task - c = Client(host=opt.host, port=opt.port) +class Client: + CHUNK = 1024 + FORMAT = pyaudio.paInt16 + CHANNELS = 1 + RATE = 16000 + RECORD_SECONDS = 60000 + START_RECORDING = False + multilingual = False + language = None + task = "transcribe" + + 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" + if translate: + Client.task = "translate" - # while loop to wait for server to be ready - print("Waiting for server ready ...") - while not START_RECORDING: - pass - print("Server Ready!") + 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) + + # 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) + else: + print("[ERROR]: No host or port specified.") + return + + # start websocket client in a thread + self.ws_thread = threading.Thread(target=self.client_socket.run_forever) + self.ws_thread.setDaemon(True) + self.ws_thread.start() + + self.frames = b"" + print("[INFO]: * recording") + + @staticmethod + def on_message(ws, message): + message = json.loads(message) + if message == "SERVER_READY": + Client.START_RECORDING = True + return + + if isinstance(message, dict): + Client.language = message.get("language") + lang_prob = message.get("language_prob") + print(f"[INFO]: Server detected language {Client.language} with probability {lang_prob}") + return + + 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 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') + else: + os.system('clear') + for element in word_list: + print(element) + + @staticmethod + def on_error(ws, error): + print(error) + + @staticmethod + def on_close(ws, close_status_code, close_msg): + print(f"[INFO]: Websocket connection closed.") + + @staticmethod + def on_open(ws): + print(Client.multilingual, Client.language, Client.task) + + print("[INFO]: Opened connection") + ws.send(json.dumps({ + 'multilingual': Client.multilingual, + 'language': Client.language, + 'task': Client.task + })) + + @staticmethod + def bytes_to_float_array(audio_bytes): + 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(), + input=True, + output=True, + frames_per_buffer=self.CHUNK) + try: + while True: + data = self.wf.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) + + self.wf.close() + self.stream.close() + + except KeyboardInterrupt: + self.wf.close() + self.stream.stop_stream() + self.stream.close() + self.p.terminate() + self.close_websocket() + print("[INFO]: Keyboard interrupt.") + + def close_websocket(self): + try: + self.client_socket.close() # Close the WebSocket connection + except Exception as e: + print("[ERROR]: Error closing WebSocket:", e) + + try: + self.ws_thread.join() # Wait for the WebSocket thread to finish + except Exception as e: + print("[ERROR:] Error joining WebSocket thread:", e) + + 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() + + def record(self, out_file="output_recording.wav"): + n_audio_file = 0 + # create dir for saving audio chunks + 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)): + 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: + t = threading.Thread( + target=self.write_audio_frames_to_file, + args=(self.frames[:], f"chunks/{n_audio_file}.wav", ) + ) + t.start() + n_audio_file += 1 + self.frames = b"" + + except KeyboardInterrupt: + if len(self.frames): + self.write_audio_frames_to_file( + self.frames[:], f"chunks/{n_audio_file}.wav") + n_audio_file += 1 + self.stream.stop_stream() + self.stream.close() + self.p.terminate() + self.close_websocket() + + # 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() + + +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.START_RECORDING: + pass + print("[INFO]: Server Ready!") + if audio is not None: + resampled_file = resample(audio) + self.client.play_file(resampled_file) + else: + self.client.record() - if opt.audio is not None: - resampled_file = resample(opt.audio) - c.play_file(resampled_file) - else: - c.record() diff --git a/server.py b/whisper_live/server.py similarity index 86% rename from server.py rename to whisper_live/server.py index d675ef1..f368bd2 100644 --- a/server.py +++ b/whisper_live/server.py @@ -10,48 +10,70 @@ logging.basicConfig(level = logging.INFO) from collections import deque from dataclasses import dataclass +from websockets.sync.server import serve import torch import numpy as np -from websockets.sync import server -from websockets.sync.server import serve -from transcriber import WhisperModel +from whisper_live.transcriber import WhisperModel -clients = {} -SERVER_READY = "SERVER_READY" - -def recv_audio(websocket): +class TranscriptionServer: """ - Receive audio chunks from client in an infinite loop. - """ - global clients - options = websocket.recv() - options = json.loads(options) - client = ServeClient( - websocket, - multilingual=options["multilingual"], - language=options["language"], - task=options["task"] - ) + Represents a transcription server that handles incoming audio from clients. - clients[websocket] = client - - while True: - try: - frame_data = websocket.recv() - frame_np = np.frombuffer(frame_data, np.float32) - clients[websocket].add_frames(frame_np) - - except Exception as e: - clients[websocket].cleanup() - clients.pop(websocket) - logging.info("Connection Closed.") - break + Attributes: + clients (dict): A dictionary to store connected clients. + """ + + def __init__(self): + self.clients = {} + + def recv_audio(self, websocket): + """ + Receive audio chunks from a client in an infinite loop. + + Args: + websocket (WebSocket): The WebSocket connection for the client. + """ + options = websocket.recv() + options = json.loads(options) + client = ServeClient( + websocket, + multilingual=options["multilingual"], + language=options["language"], + task=options["task"], + ) + + self.clients[websocket] = client + + while True: + try: + frame_data = websocket.recv() + frame_np = np.frombuffer(frame_data, np.float32) + self.clients[websocket].add_frames(frame_np) + + except Exception as e: + self.clients[websocket].cleanup() + self.clients.pop(websocket) + logging.info("Connection Closed.") + break + + def run(self, host, port): + """ + Run the transcription server. + + Args: + host (str): The host address to bind the server. + port (int): The port number to bind the server. + """ + with serve(self.recv_audio, host, port) as server: + server.serve_forever() class ServeClient: RATE = 16000 + SERVER_READY = "SERVER_READY" + def __init__(self, websocket, task="transcribe", device=None, multilingual=False, language=None): self.data = b"" self.frames = b"" @@ -94,7 +116,7 @@ class ServeClient: self.websocket = websocket self.trans_thread = threading.Thread(target=self.speech_to_text) self.trans_thread.start() - self.websocket.send(json.dumps(SERVER_READY)) + self.websocket.send(json.dumps(self.SERVER_READY)) def fill_output(self, output): """ @@ -302,8 +324,4 @@ class ServeClient: logging.info("Cleaning up.") self.exit = True self.transcriber.destroy() - -if __name__ == "__main__": - with serve(recv_audio, "0.0.0.0", 9090) as server: - server.serve_forever() \ No newline at end of file diff --git a/transcriber.py b/whisper_live/transcriber.py similarity index 100% rename from transcriber.py rename to whisper_live/transcriber.py