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

Replace ffmpeg with av lib for resampling, rtsp & hls streams

Browse Source 
Signed-off-by: makaveli10 <vineet.suryan@collabora.com>
This commit is contained in:
makaveli10
2025-01-22 05:11:04 -05:00
parent 893265bb3f
commit 5aa5826f36
9 changed files with 110 additions and 102 deletions
Download Patch File Download Diff File Expand all files Collapse all files
whisper_live/client.py
+57 -58
View File
@@ -10,7 +10,7 @@ import json
import websocket
import uuid
import time
import ffmpeg
import av
import whisper_live.utils as utils
@@ -421,84 +421,83 @@ class TranscriptionTeeClient:
def process_rtsp_stream(self, rtsp_url):
"""
Connect to an RTSP source, process the audio stream, and send it for trascription.
Connect to an RTSP source, process the audio stream, and send it for transcription.
Args:
rtsp_url (str): The URL of the RTSP stream source.
"""
process = self.get_rtsp_ffmpeg_process(rtsp_url)
self.handle_ffmpeg_process(process, stream_type='RTSP')
print("[INFO]: Connecting to RTSP stream...")
try:
container = av.open(rtsp_url, format="rtsp", options={"rtsp_transport": "tcp"})
self.process_av_stream(container, stream_type="RTSP")
except Exception as e:
print(f"[ERROR]: Failed to process RTSP stream: {e}")
finally:
for client in self.clients:
client.wait_before_disconnect()
self.multicast_packet(Client.END_OF_AUDIO.encode('utf-8'), True)
self.close_all_clients()
self.write_all_clients_srt()
print("[INFO]: RTSP stream processing finished.")
def process_hls_stream(self, hls_url, save_file):
def process_hls_stream(self, hls_url, save_file=None):
"""
Connect to an HLS source, process the audio stream, and send it for transcription.
Args:
hls_url (str): The URL of the HLS stream source.
save_file （str, optional): Local path to save the network stream.
save_file (str, optional): Local path to save the network stream.
"""
process = self.get_hls_ffmpeg_process(hls_url, save_file)
self.handle_ffmpeg_process(process, stream_type='HLS')
def handle_ffmpeg_process(self, process, stream_type):
print(f"[INFO]: Connecting to {stream_type} stream...")
stderr_thread = threading.Thread(target=self.consume_stderr, args=(process,))
stderr_thread.start()
print("[INFO]: Connecting to HLS stream...")
try:
# Process the stream
while True:
in_bytes = process.stdout.read(self.chunk * 2) # 2 bytes per sample
if not in_bytes:
break
audio_array = self.bytes_to_float_array(in_bytes)
self.multicast_packet(audio_array.tobytes())
container = av.open(hls_url, format="hls")
self.process_av_stream(container, stream_type="HLS", save_file=save_file)
except Exception as e:
print(f"[ERROR]: Failed to connect to {stream_type} stream: {e}")
print(f"[ERROR]: Failed to process HLS stream: {e}")
finally:
for client in self.clients:
client.wait_before_disconnect()
self.multicast_packet(Client.END_OF_AUDIO.encode('utf-8'), True)
self.close_all_clients()
self.write_all_clients_srt()
if process:
process.kill()
print("[INFO]: HLS stream processing finished.")
print(f"[INFO]: {stream_type} stream processing finished.")
def get_rtsp_ffmpeg_process(self, rtsp_url):
return (
ffmpeg
.input(rtsp_url, threads=0)
.output('-', format='s16le', acodec='pcm_s16le', ac=1, ar=self.rate)
.run_async(pipe_stdout=True, pipe_stderr=True)
)
def get_hls_ffmpeg_process(self, hls_url, save_file):
if save_file is None:
process = (
ffmpeg
.input(hls_url, threads=0)
.output('-', format='s16le', acodec='pcm_s16le', ac=1, ar=self.rate)
.run_async(pipe_stdout=True, pipe_stderr=True)
)
else:
input = ffmpeg.input(hls_url, threads=0)
output_file = input.output(save_file, acodec='copy', vcodec='copy').global_args('-loglevel', 'quiet')
output_std = input.output('-', format='s16le', acodec='pcm_s16le', ac=1, ar=self.rate)
process = (
ffmpeg.merge_outputs(output_file, output_std)
.run_async(pipe_stdout=True, pipe_stderr=True)
)
return process
def consume_stderr(self, process):
def process_av_stream(self, container, stream_type, save_file=None):
"""
Consume and log the stderr output of a process in a separate thread.
Process an AV container stream and send audio packets to the server.
Args:
process (subprocess.Popen): The process whose stderr output will be logged.
container (av.container.InputContainer): The input container to process.
stream_type (str): The type of stream being processed ("RTSP" or "HLS").
save_file (str, optional): Local path to save the stream. Default is None.
"""
for line in iter(process.stderr.readline, b""):
logging.debug(f'[STDERR]: {line.decode()}')
audio_stream = next((s for s in container.streams if s.type == "audio"), None)
if not audio_stream:
print(f"[ERROR]: No audio stream found in {stream_type} source.")
return
output_container = None
if save_file:
output_container = av.open(save_file, mode="w")
output_audio_stream = output_container.add_stream(codec_name="pcm_s16le", rate=self.rate)
try:
for packet in container.demux(audio_stream):
for frame in packet.decode():
audio_data = frame.to_ndarray().tobytes()
self.multicast_packet(audio_data)
if save_file:
output_container.mux(frame)
except Exception as e:
print(f"[ERROR]: Error during {stream_type} stream processing: {e}")
finally:
# Wait for server to send any leftover transcription.
time.sleep(5)
self.multicast_packet(Client.END_OF_AUDIO.encode('utf-8'), True)
if output_container:
output_container.close()
container.close()
def save_chunk(self, n_audio_file):
"""