diff --git a/Dockerfile b/Dockerfile index 6193dac..8805246 100644 --- a/Dockerfile +++ b/Dockerfile @@ -39,6 +39,7 @@ COPY requirements/ /app RUN bash setup.sh RUN pip install -r server.txt -COPY *py /app +RUN pip install whisper-live -CMD ["python", "server.py"] +COPY run_server.py /app +CMD ["python", "run_server.py"] diff --git a/README.md b/README.md index dba34fd..f9fd9cf 100644 --- a/README.md +++ b/README.md @@ -11,41 +11,43 @@ Unlike traditional speech recognition systems that rely on continuous audio stre bash setup.sh ``` -- To install client requirements +- Install whisper-live from pip ```bash - pip install -r requirements/client.txt -``` - -- To install server requirements -```bash - pip install -r requirements/server.txt + pip install whisper-live ``` ## Getting Started - Run the server -```bash - python server.py +```python + from whisper_live.server import TranscriptionServer + server = TranscriptionServer() + server.run("0.0.0.0", 9090) ``` - On the client side - To transcribe an audio file: - ```bash - python client.py --audio "audio.wav" --host "localhost" --port "9090" --multilingual --language "hi" --task "transcribe" - "translate" + ```python + from whisper_live.client import TranscriptionClient + client = TranscriptionClient("localhost", 9090, multilingual=True, language="hi", translate=True) + client(audio_file_path) ``` - This command transcribes the specified audio file (audio.wav) using the Whisper model. It connects to the server running on localhost at port 9090. It also enables the multilingual feature, allowing transcription in multiple languages. The --language flag specifies the target language for transcription, in this case, Hindi ("hi"). The --task flag is set to "transcribe" to indicate that transcription is the desired task. Also, --task can be set to "translate" to translate source language to English. + This command transcribes the specified audio file (audio.wav) using the Whisper model. It connects to the server running on localhost at port 9090. It also enables the multilingual feature, allowing transcription in multiple languages. The language option specifies the target language for transcription, in this case, Hindi ("hi"). The translate option should be set to `True` if we want to translate from the source language to English and `False` if we want to transcribe in the source language. - To transcribe from microphone: - ```bash - python client.py --host "localhost" --port "9090" --multilingual --language "en" --task "transcribe" + ```python + from whisper_live.client import TranscriptionClient + client = TranscriptionClient(host, port, multilingual=True, language="hi", translate=True) + client() ``` This command captures audio from the microphone and sends it to the server for transcription. It uses the same options as the previous command, enabling the multilingual feature and specifying the target language and task. ## Transcribe audio from browser - Run the server -```bash - python server.py +```python + from whisper_live.server import TranscriptionServer + server = TranscriptionServer() + server.run("0.0.0.0", 9090) ``` This would start the websocket server on port ```9090```. @@ -67,9 +69,8 @@ This would start the websocket server on port ```9090```. ``` ## Future Work -- [x] Update Documentation. -- [x] Keep only a single server implementation i.e. websockets and get rid of the socket implementation in ```server.py```. Also, update ```client.py``` to websockets-client implemenation. - [ ] Add translation to other languages on top of transcription. +- [ ] TensorRT backend for Whisper. ## Citations ```bibtex diff --git a/requirements/client.txt b/requirements/client.txt index 3befca2..359cc04 100644 --- a/requirements/client.txt +++ b/requirements/client.txt @@ -1,5 +1,4 @@ PyAudio ffmpeg-python scipy -websocket-client -onnxruntime \ No newline at end of file +websocket-client \ No newline at end of file diff --git a/requirements/server.txt b/requirements/server.txt index 2a34780..360e392 100644 --- a/requirements/server.txt +++ b/requirements/server.txt @@ -3,4 +3,5 @@ faster-whisper==0.6.0 --extra-index-url https://download.pytorch.org/whl/cu111 torch==1.10.1 torchaudio==0.10.1 -websockets \ No newline at end of file +websockets +onnxruntime \ No newline at end of file diff --git a/run_server.py b/run_server.py new file mode 100644 index 0000000..9cf6165 --- /dev/null +++ b/run_server.py @@ -0,0 +1,6 @@ + +from whisper_live.server import TranscriptionServer + +if __name__ == "__main__": + server = TranscriptionServer() + server.run_server("0.0.0.0", 9090) \ No newline at end of file diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..de5f007 --- /dev/null +++ b/setup.py @@ -0,0 +1,44 @@ +import pathlib +from setuptools import find_packages, setup + +# The directory containing this file +HERE = pathlib.Path(__file__).parent + +# The text of the README file +README = (HERE / "README.md").read_text() + +# This call to setup() does all the work +setup(name="whisper-live", + version="0.0.4", + description="A nearly-live implementation of OpenAI's Whisper.", + long_description=README, + long_description_content_type="text/markdown", + include_package_data=True, + url="https://github.com/collabora/WhisperLive", + author="Collabora Ltd", + author_email="vineet.suryan@collabora.com", + license="MIT", + classifiers=[ + "License :: OSI Approved :: MIT License", + "Programming Language :: Python :: 3.8", + "Programming Language :: Python :: 3.9", + ], + packages=find_packages( + exclude=("examples", + "Audio-Transcription-Chrome", + "Audio-Transcription-Firefox", + "requirements", + "whisper-finetuning" + ) + ), + install_requires=[ + "PyAudio", + "faster-whisper==0.6.0", + "torch", + "torchaudio", + "websockets", + "onnxruntime", + "ffmpeg-python", + "scipy", + "websocket-client", + ]) \ No newline at end of file diff --git a/setup.sh b/setup.sh index 1ef9996..13a015c 100644 --- a/setup.sh +++ b/setup.sh @@ -1,3 +1,3 @@ #! /bin/bash -apt-get install portaudio19-dev -y +apt-get install portaudio19-dev ffmpeg -y 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