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 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
INSTANCES = {}
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"
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)
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)
@@ -93,30 +103,34 @@ 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():
@@ -126,48 +140,47 @@ class Client:
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):
@@ -178,36 +191,41 @@ class Client:
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:
@@ -224,12 +242,12 @@ class Client:
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,9 +255,10 @@ 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)
@@ -247,10 +266,13 @@ class Client:
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()
@@ -270,21 +293,26 @@ class Client:
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:
@@ -293,8 +321,8 @@ class TranscriptionClient:
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()