cjones/WhisperLive
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Tidy up the client python file

Browse Source
This commit is contained in:
Ryan Pavlik
2023-09-22 14:30:52 -05:00
parent d319aa2fc9
commit a65d69cfbd
1 changed files with 158 additions and 131 deletions
Download Patch File Download Diff File Expand all files Collapse all files
whisper_live/client.py
+146 -119
View File
@@ -1,5 +1,4 @@
import os import os
import argparse
import wave import wave
import numpy as np import numpy as np
@@ -44,47 +43,58 @@ def resample(file: str, sr: int = 16000):
class Client: class Client:
CHUNK = 1024 INSTANCES = {}
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): def __init__(
Client.multilingual = is_multilingual self, host=None, port=None, is_multilingual=False, lang=None, translate=False
Client.language = lang if is_multilingual else "en" ):
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: if translate:
Client.task = "translate" self.task = "translate"
self.timestamp_offset = 0.0 self.timestamp_offset = 0.0
self.audio_bytes = None self.audio_bytes = None
self.p = pyaudio.PyAudio() self.p = pyaudio.PyAudio()
self.stream = self.p.open(format=self.FORMAT, self.stream = self.p.open(
channels=self.CHANNELS, format=self.format,
rate=self.RATE, channels=self.channels,
input=True, rate=self.rate,
frames_per_buffer=self.CHUNK) input=True,
frames_per_buffer=self.chunk,
)
# create websocket connection # create websocket connection
if host is not None and port is not None: if host is not None and port is not None:
socket_url = f"ws://{host}:{port}" socket_url = f"ws://{host}:{port}"
self.client_socket = websocket.WebSocketApp(socket_url, self.client_socket = websocket.WebSocketApp(
on_open=Client.on_open, socket_url,
on_message=Client.on_message, on_open=lambda ws: self.on_open(ws),
on_error=Client.on_error, on_message=lambda ws, message: self.on_message(ws, message),
on_close=Client.on_close) 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: else:
print("[ERROR]: No host or port specified.") print("[ERROR]: No host or port specified.")
return return
Client.INSTANCES[self.uid] = self
# start websocket client in a thread # start websocket client in a thread
self.ws_thread = threading.Thread(target=self.client_socket.run_forever) self.ws_thread = threading.Thread(target=self.client_socket.run_forever)
self.ws_thread.setDaemon(True) self.ws_thread.setDaemon(True)
@@ -93,30 +103,34 @@ class Client:
self.frames = b"" self.frames = b""
print("[INFO]: * recording") print("[INFO]: * recording")
@staticmethod def on_message(self, ws, message):
def on_message(ws, message): self.last_response_recieved = time.time()
Client.LAST_RESPONSE_RECIEVED = time.time()
message = json.loads(message) message = json.loads(message)
if message.get('uid')!=Client.uid:
if self.uid != message.get("uid"):
print("[ERROR]: invalid client uid") print("[ERROR]: invalid client uid")
return return
if "status" in message.keys() and message["status"] == "WAIT": if "status" in message.keys() and message["status"] == "WAIT":
Client.WAITING = True self.waiting = True
print(f"[INFO]:Server is full. Estimated wait time {round(message['message'])} minutes.") print(
f"[INFO]:Server is full. Estimated wait time {round(message['message'])} minutes."
)
if "message" in message.keys() and message["message"] == "DISCONNECT": if "message" in message.keys() and message["message"] == "DISCONNECT":
print("[INFO]: Server overtime disconnected.") print("[INFO]: Server overtime disconnected.")
Client.RECORDING = False self.recording = False
if "message" in message.keys() and message["message"] == "SERVER_READY": if "message" in message.keys() and message["message"] == "SERVER_READY":
Client.RECORDING = True self.recording = True
return return
if "language" in message.keys(): if "language" in message.keys():
Client.language = message.get("language") self.language = message.get("language")
lang_prob = message.get("language_prob") 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 return
if "segments" not in message.keys(): if "segments" not in message.keys():
@@ -126,48 +140,47 @@ class Client:
text = [] text = []
if len(message): if len(message):
for seg in message: for seg in message:
if len(text): if text and text[-1] == seg["text"]:
if text[-1] != seg["text"]: # already got it
text.append(seg["text"]) continue
else: text.append(seg["text"])
text.append(seg["text"]) # keep only last 3
if len(text) > 3: if len(text) > 3:
text = text[-3:] text = text[-3:]
wrapper = textwrap.TextWrapper(width=60) wrapper = textwrap.TextWrapper(width=60)
word_list = wrapper.wrap(text="".join(text)) word_list = wrapper.wrap(text="".join(text))
# Print each line. # Print each line.
if os.name=='nt': if os.name == "nt":
os.system('cls') os.system("cls")
else: else:
os.system('clear') os.system("clear")
for element in word_list: for element in word_list:
print(element) print(element)
@staticmethod def on_error(self, ws, error):
def on_error(ws, error):
print(error) print(error)
@staticmethod def on_close(self, ws, close_status_code, close_msg):
def on_close(ws, close_status_code, close_msg): print(f"[INFO]: Websocket connection closed: {close_status_code}: {close_msg}")
print(f"[INFO]: Websocket connection closed.")
@staticmethod def on_open(self, ws):
def on_open(ws): print(self.multilingual, self.language, self.task)
print(Client.multilingual, Client.language, Client.task)
print("[INFO]: Opened connection") print("[INFO]: Opened connection")
ws.send(json.dumps({ ws.send(
'uid': Client.uid, json.dumps(
'multilingual': Client.multilingual, {
'language': Client.language, "uid": self.uid,
'task': Client.task "multilingual": self.multilingual,
})) "language": self.language,
"task": self.task,
}
)
)
@staticmethod @staticmethod
def bytes_to_float_array(audio_bytes): def bytes_to_float_array(audio_bytes):
raw_data = np.frombuffer( raw_data = np.frombuffer(buffer=audio_bytes, dtype=np.int16)
buffer=audio_bytes, dtype=np.int16
)
return raw_data.astype(np.float32) / 32768.0 return raw_data.astype(np.float32) / 32768.0
def send_packet_to_server(self, message): def send_packet_to_server(self, message):
@@ -178,36 +191,41 @@ class Client:
def play_file(self, filename): def play_file(self, filename):
# read audio and create pyaudio stream # read audio and create pyaudio stream
self.wf = wave.open(filename, 'rb') with wave.open(filename, "rb") as wavfile:
self.stream = self.p.open(format=self.p.get_format_from_width(self.wf.getsampwidth()), self.stream = self.p.open(
channels=self.wf.getnchannels(), format=self.p.get_format_from_width(wavfile.getsampwidth()),
rate=self.wf.getframerate(), channels=wavfile.getnchannels(),
rate=wavfile.getframerate(),
input=True, input=True,
output=True, output=True,
frames_per_buffer=self.CHUNK) frames_per_buffer=self.chunk,
try: )
while Client.RECORDING: try:
data = self.wf.readframes(self.CHUNK) while self.recording:
if data==b'': break data = wavfile.readframes(self.chunk)
if data == b"":
break
audio_array = Client.bytes_to_float_array(data) audio_array = self.bytes_to_float_array(data)
self.send_packet_to_server(audio_array.tobytes()) self.send_packet_to_server(audio_array.tobytes())
self.stream.write(data) self.stream.write(data)
self.wf.close() wavfile.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()
except KeyboardInterrupt: assert self.last_response_recieved
self.wf.close() elapsed_time = time.time() - self.last_response_recieved
self.stream.stop_stream() while elapsed_time < self.disconnect_if_no_response_for:
self.stream.close() continue
self.p.terminate() self.stream.close()
self.close_websocket() self.close_websocket()
print("[INFO]: Keyboard interrupt.")
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): def close_websocket(self):
try: try:
@@ -224,12 +242,12 @@ class Client:
return self.client_socket return self.client_socket
def write_audio_frames_to_file(self, frames, file_name): def write_audio_frames_to_file(self, frames, file_name):
wf = wave.open(file_name, 'wb') with wave.open(file_name, "wb") as wavfile:
wf.setnchannels(self.CHANNELS) wavfile: wave.Wave_write
wf.setsampwidth(2) wavfile.setnchannels(self.channels)
wf.setframerate(self.RATE) wavfile.setsampwidth(2)
wf.writeframes(frames) wavfile.setframerate(self.rate)
wf.close() wavfile.writeframes(frames)
def record(self, out_file="output_recording.wav"): def record(self, out_file="output_recording.wav"):
n_audio_file = 0 n_audio_file = 0
@@ -237,9 +255,10 @@ class Client:
if not os.path.exists("chunks"): if not os.path.exists("chunks"):
os.makedirs("chunks", exist_ok=True) os.makedirs("chunks", exist_ok=True)
try: try:
for _ in range(0, int(self.RATE / self.CHUNK * self.RECORD_SECONDS)): for _ in range(0, int(self.rate / self.chunk * self.record_seconds)):
if not Client.RECORDING: break if not self.recording:
data = self.stream.read(self.CHUNK) break
data = self.stream.read(self.chunk)
self.frames += data self.frames += data
audio_array = Client.bytes_to_float_array(data) audio_array = Client.bytes_to_float_array(data)
@@ -247,10 +266,13 @@ class Client:
self.send_packet_to_server(audio_array.tobytes()) self.send_packet_to_server(audio_array.tobytes())
# save frames if more than a minute # save frames if more than a minute
if len(self.frames) > 60*self.RATE: if len(self.frames) > 60 * self.rate:
t = threading.Thread( t = threading.Thread(
target=self.write_audio_frames_to_file, 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() t.start()
n_audio_file += 1 n_audio_file += 1
@@ -259,7 +281,8 @@ class Client:
except KeyboardInterrupt: except KeyboardInterrupt:
if len(self.frames): if len(self.frames):
self.write_audio_frames_to_file( 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 n_audio_file += 1
self.stream.stop_stream() self.stream.stop_stream()
self.stream.close() self.stream.close()
@@ -270,21 +293,26 @@ class Client:
self.write_output_recording(n_audio_file, out_file) self.write_output_recording(n_audio_file, out_file)
def write_output_recording(self, 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")] input_files = [
wf = wave.open(out_file, 'wb') f"chunks/{i}.wav"
wf.setnchannels(self.CHANNELS) for i in range(n_audio_file)
wf.setsampwidth(2) if os.path.exists(f"chunks/{i}.wav")
wf.setframerate(self.RATE) ]
for in_file in input_files: with wave.open(out_file, "wb") as wavfile:
w = wave.open(in_file, 'rb') wavfile: wave.Wave_write
while True: wavfile.setnchannels(self.channels)
data = w.readframes(self.CHUNK) wavfile.setsampwidth(2)
if data==b'': break wavfile.setframerate(self.rate)
wf.writeframes(data) for in_file in input_files:
w.close() with wave.open(in_file, "rb") as wav_in:
# remove this file while True:
os.remove(in_file) data = wav_in.readframes(self.chunk)
wf.close() if data == b"":
break
wavfile.writeframes(data)
# remove this file
os.remove(in_file)
wavfile.close()
class TranscriptionClient: class TranscriptionClient:
@@ -293,8 +321,8 @@ class TranscriptionClient:
def __call__(self, audio=None): def __call__(self, audio=None):
print("[INFO]: Waiting for server ready ...") print("[INFO]: Waiting for server ready ...")
while not Client.RECORDING: while not self.client.recording:
if Client.WAITING: if self.client.waiting:
self.client.close_websocket() self.client.close_websocket()
return return
pass pass
@@ -304,4 +332,3 @@ class TranscriptionClient:
self.client.play_file(resampled_file) self.client.play_file(resampled_file)
else: else:
self.client.record() self.client.record()