Ryan Pavlik
335 lines
11 KiB
Python
335 lines
11 KiB
Python
import os
|
|
import wave
|
|
|
|
import numpy as np
|
|
import scipy
|
|
import ffmpeg
|
|
import pyaudio
|
|
import threading
|
|
import textwrap
|
|
import json
|
|
import websocket
|
|
import uuid
|
|
import time
|
|
|
|
|
|
def resample(file: str, sr: int = 16000):
|
|
"""
|
|
# https://github.com/openai/whisper/blob/7858aa9c08d98f75575035ecd6481f462d66ca27/whisper/audio.py#L22
|
|
Open an audio file and read as mono waveform, resampling as necessary,
|
|
save the resampled audio
|
|
Parameters
|
|
----------
|
|
file: str
|
|
The audio file to open
|
|
sr: int
|
|
The sample rate to resample the audio if necessary
|
|
"""
|
|
try:
|
|
# This launches a subprocess to decode audio while down-mixing and resampling as necessary.
|
|
# Requires the ffmpeg CLI and `ffmpeg-python` package to be installed.
|
|
out, _ = (
|
|
ffmpeg.input(file, threads=0)
|
|
.output("-", format="s16le", acodec="pcm_s16le", ac=1, ar=sr)
|
|
.run(cmd=["ffmpeg", "-nostdin"], capture_stdout=True, capture_stderr=True)
|
|
)
|
|
except ffmpeg.Error as e:
|
|
raise RuntimeError(f"Failed to load audio: {e.stderr.decode()}") from e
|
|
np_buffer = np.frombuffer(out, dtype=np.int16)
|
|
|
|
resampled_file = f"{file.split('.')[0]}_resampled.wav"
|
|
scipy.io.wavfile.write(resampled_file, sr, np_buffer.astype(np.int16))
|
|
return resampled_file
|
|
|
|
|
|
class Client:
|
|
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:
|
|
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,
|
|
)
|
|
|
|
# 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=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)
|
|
self.ws_thread.start()
|
|
|
|
self.frames = b""
|
|
print("[INFO]: * recording")
|
|
|
|
def on_message(self, ws, message):
|
|
self.last_response_recieved = time.time()
|
|
message = json.loads(message)
|
|
|
|
if self.uid != message.get("uid"):
|
|
print("[ERROR]: invalid client uid")
|
|
return
|
|
|
|
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.")
|
|
self.recording = False
|
|
|
|
if "message" in message.keys() and message["message"] == "SERVER_READY":
|
|
self.recording = True
|
|
return
|
|
|
|
if "language" in message.keys():
|
|
self.language = message.get("language")
|
|
lang_prob = message.get("language_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 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")
|
|
else:
|
|
os.system("clear")
|
|
for element in word_list:
|
|
print(element)
|
|
|
|
def on_error(self, ws, error):
|
|
print(error)
|
|
|
|
def on_close(self, ws, close_status_code, close_msg):
|
|
print(f"[INFO]: Websocket connection closed: {close_status_code}: {close_msg}")
|
|
|
|
def on_open(self, ws):
|
|
print(self.multilingual, self.language, self.task)
|
|
|
|
print("[INFO]: Opened connection")
|
|
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)
|
|
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
|
|
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 self.recording:
|
|
data = wavfile.readframes(self.chunk)
|
|
if data == b"":
|
|
break
|
|
|
|
audio_array = self.bytes_to_float_array(data)
|
|
self.send_packet_to_server(audio_array.tobytes())
|
|
self.stream.write(data)
|
|
|
|
wavfile.close()
|
|
|
|
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:
|
|
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):
|
|
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
|
|
# 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)):
|
|
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:
|
|
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")
|
|
]
|
|
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 self.client.recording:
|
|
if self.client.waiting:
|
|
self.client.close_websocket()
|
|
return
|
|
pass
|
|
print("[INFO]: Server Ready!")
|
|
if audio is not None:
|
|
resampled_file = resample(audio)
|
|
self.client.play_file(resampled_file)
|
|
else:
|
|
self.client.record()
|