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 argparse
import wave
import numpy as np
@@ -45,44 +44,11 @@ def resample(file: str, sr: int = 16000):
class Client:
"""
Represents a client for audio recording and streaming to a server using WebSocket communication.
INSTANCES = {}
This class allows audio recording from the microphone or playing audio from a file while streaming it to
a server for transcription or translation. It uses PyAudio for audio recording and playback and WebSocket
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):
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.
@@ -96,42 +62,54 @@ class Client:
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.
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
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:
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,
)
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)
self.ws_thread.start()
@@ -139,8 +117,7 @@ class Client:
self.frames = b""
print("[INFO]: * recording")
@staticmethod
def on_message(ws, message):
def on_message(self, ws, message):
"""
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.
"""
Client.LAST_RESPONSE_RECIEVED = time.time()
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():
@@ -184,33 +166,30 @@ 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):
def on_open(self, ws):
"""
Callback function called when the WebSocket connection is successfully opened.
@@ -221,15 +200,19 @@ class Client:
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")
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):
@@ -245,9 +228,7 @@ class Client:
Returns:
np.ndarray: A NumPy array containing the audio data as float values normalized between -1 and 1.
"""
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):
@@ -277,36 +258,42 @@ class Client:
Args:
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()),
channels=self.wf.getnchannels(),
rate=self.wf.getframerate(),
# 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 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
while time.time() - self.last_response_recieved < 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):
"""
@@ -347,12 +334,12 @@ class Client:
file_name (str): The name of the WAV file to which the frames will be written.
"""
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"):
"""
@@ -375,9 +362,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)
@@ -385,10 +373,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
@@ -397,7 +388,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()
@@ -420,21 +412,26 @@ class Client:
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")]
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:
@@ -477,8 +474,8 @@ class TranscriptionClient:
"""
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
@@ -488,4 +485,3 @@ class TranscriptionClient:
self.client.play_file(resampled_file)
else:
self.client.record()