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

merge upstream

Browse Source
This commit is contained in:
makaveli10
2023-09-25 16:14:33 +05:30
parent 9288791943 e89afcc40f
commit c9f5d48c05
1 changed files with 159 additions and 163 deletions
Download Patch File Download Diff File Expand all files Collapse all files
whisper_live/client.py
+148 -152
View File
@@ -1,5 +1,4 @@
import os import os
import argparse
import wave import wave
import numpy as np import numpy as np
@@ -45,44 +44,11 @@ def resample(file: str, sr: int = 16000):
class Client: class Client:
""" INSTANCES = {}
Represents a client for audio recording and streaming to a server using WebSocket communication.
This class allows audio recording from the microphone or playing audio from a file while streaming it to def __init__(
a server for transcription or translation. It uses PyAudio for audio recording and playback and WebSocket self, host=None, port=None, is_multilingual=False, lang=None, translate=False
for communication with the server. ):
Attributes:
CHUNK (int): The size of audio chunks for recording and playback.
FORMAT: The audio format used by PyAudio (paInt16 for 16-bit PCM).
CHANNELS (int): The number of audio channels (1 for mono).
RATE (int): The audio sampling rate in Hz (samples per second).
RECORD_SECONDS (int): The maximum duration for audio recording in seconds.
RECORDING (bool): Indicates whether recording is currently active.
multilingual (bool): Indicates if multilingual transcription is enabled.
language (str): The selected language for transcription.
task (str): The transcription or translation task to be performed.
uid (str): A unique identifier for the client.
WAITING (bool): Indicates if the client is waiting for server availability.
LAST_RESPONSE_RECIEVED (float): Timestamp of the last response received from the server.
DISCONNECT_IF_NO_RESPONSE_FOR (int): Maximum time without server response before disconnection.
"""
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
def __init__(self, host=None, port=None, is_multilingual=False, lang=None, translate=False):
""" """
Initializes a Client instance for audio recording and streaming to a server. Initializes a Client instance for audio recording and streaming to a server.
@@ -96,42 +62,54 @@ class Client:
is_multilingual (bool, optional): Specifies if multilingual transcription is enabled. Default is False. is_multilingual (bool, optional): Specifies if multilingual transcription is enabled. Default is False.
lang (str, optional): The selected language for transcription when multilingual is disabled. Default is None. lang (str, optional): The selected language for transcription when multilingual is disabled. Default is None.
translate (bool, optional): Specifies if the task is translation. Default is False. translate (bool, optional): Specifies if the task is translation. Default is False.
Attributes:
timestamp_offset (float): A timestamp offset for tracking audio timing.
audio_bytes (bytes): A buffer for storing audio data.
p (pyaudio.PyAudio): An instance of PyAudio for audio streaming.
stream (pyaudio.Stream): The audio stream for recording.
client_socket (websocket.WebSocketApp): The WebSocket client for server communication.
ws_thread (threading.Thread): A thread for running the WebSocket client.
frames (bytes): A buffer for accumulating audio frames.
""" """
Client.multilingual = is_multilingual self.chunk = 1024
Client.language = lang if is_multilingual else "en" 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,
)
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
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)
self.ws_thread.start() self.ws_thread.start()
@@ -139,8 +117,7 @@ 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):
""" """
Callback function called when a message is received from the server. Callback function called when a message is received from the server.
@@ -153,28 +130,33 @@ class Client:
message (str): The received message from the server. message (str): The received message from the server.
""" """
Client.LAST_RESPONSE_RECIEVED = time.time() self.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():
@@ -184,33 +166,30 @@ 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):
""" """
Callback function called when the WebSocket connection is successfully opened. Callback function called when the WebSocket connection is successfully opened.
@@ -221,15 +200,19 @@ class Client:
ws (websocket.WebSocketApp): The WebSocket client instance. ws (websocket.WebSocketApp): The WebSocket client instance.
""" """
print(Client.multilingual, Client.language, Client.task) print(self.multilingual, self.language, self.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):
@@ -245,9 +228,7 @@ class Client:
Returns: Returns:
np.ndarray: A NumPy array containing the audio data as float values normalized between -1 and 1. np.ndarray: A NumPy array containing the audio data as float values normalized between -1 and 1.
""" """
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):
@@ -277,36 +258,42 @@ class Client:
Args: Args:
filename (str): The path to the audio file to be played and sent to the server. filename (str): The path to the audio file to be played and sent to the server.
""" """
self.wf = wave.open(filename, 'rb')
self.stream = self.p.open(format=self.p.get_format_from_width(self.wf.getsampwidth()), # read audio and create pyaudio stream
channels=self.wf.getnchannels(), with wave.open(filename, "rb") as wavfile:
rate=self.wf.getframerate(), self.stream = self.p.open(
format=self.p.get_format_from_width(wavfile.getsampwidth()),
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() while time.time() - self.last_response_recieved < self.disconnect_if_no_response_for:
self.stream.stop_stream() continue
self.stream.close() self.stream.close()
self.p.terminate() 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):
""" """
@@ -347,12 +334,12 @@ class Client:
file_name (str): The name of the WAV file to which the frames will be written. file_name (str): The name of the WAV file to which the frames will be written.
""" """
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"):
""" """
@@ -375,9 +362,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)
@@ -385,10 +373,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
@@ -397,7 +388,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()
@@ -420,21 +412,26 @@ class Client:
out_file (str): The name of the output WAV file to save the final recording. out_file (str): The name of the output WAV file to save the final recording.
""" """
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:
@@ -477,8 +474,8 @@ class TranscriptionClient:
""" """
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
@@ -488,4 +485,3 @@ class TranscriptionClient:
self.client.play_file(resampled_file) self.client.play_file(resampled_file)
else: else:
self.client.record() self.client.record()