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

Signed-off-by: makaveli10 <vineet.suryan@collabora.com>

whisper_live/client.py

@@ -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):
         """

whisper_live/tensorrt_utils.py

@@ -23,8 +23,12 @@ from typing import Dict, Iterable, List, Optional, TextIO, Tuple, Union
 import kaldialign
 import numpy as np
 import soundfile
+import av
+import wave
 import torch
 import torch.nn.functional as F
+from whisper_live.utils import resample
+
 
 Pathlike = Union[str, Path]
 
@@ -35,38 +39,33 @@ CHUNK_LENGTH = 30
 N_SAMPLES = CHUNK_LENGTH * SAMPLE_RATE  # 480000 samples in a 30-second chunk
 
 
-def load_audio(file: str, sr: int = SAMPLE_RATE):
+def load_audio(file: str, sr: int = 16000):
     """
-    Open an audio file and read as mono waveform, resampling as necessary
+    Open an audio file, resample it, and read as a mono waveform.
 
     Parameters
     ----------
     file: str
-        The audio file to open
+        The audio file to open.
 
     sr: int
-        The sample rate to resample the audio if necessary
+        The sample rate to resample the audio if necessary.
 
     Returns
     -------
     A NumPy array containing the audio waveform, in float32 dtype.
     """
+    resampled_file = resample(file, sr)
 
-    # This launches a subprocess to decode audio while down-mixing
-    # and resampling as necessary.  Requires the ffmpeg CLI in PATH.
-    # fmt: off
-    cmd = [
-        "ffmpeg", "-nostdin", "-threads", "0", "-i", file, "-f", "s16le", "-ac",
-        "1", "-acodec", "pcm_s16le", "-ar",
-        str(sr), "-"
-    ]
-    # fmt: on
-    try:
-        out = run(cmd, capture_output=True, check=True).stdout
-    except CalledProcessError as e:
-        raise RuntimeError(f"Failed to load audio: {e.stderr.decode()}") from e
+    with wave.open(resampled_file, "rb") as wav_file:
+        num_frames = wav_file.getnframes()
+        raw_data = wav_file.readframes(num_frames)
 
-    return np.frombuffer(out, np.int16).flatten().astype(np.float32) / 32768.0
+        audio_data = np.frombuffer(raw_data, dtype=np.int16)
+
+    audio_data = audio_data.astype(np.float32) / 32768.0
+
+    return audio_data
 
 
 def load_audio_wav_format(wav_path):

whisper_live/utils.py

@@ -1,8 +1,9 @@
 import os
 import textwrap
 import scipy
-import ffmpeg
 import numpy as np
+import av
+from pathlib import Path
 
 
 def clear_screen():
@@ -26,8 +27,8 @@ def format_time(s):
     return f"{hours:02}:{minutes:02}:{seconds:02},{milliseconds:03}"
 
 
-def create_srt_file(segments, output_file):
-    with open(output_file, 'w', encoding='utf-8') as srt_file:
+def create_srt_file(segments, resampled_file):
+    with open(resampled_file, 'w', encoding='utf-8') as srt_file:
         segment_number = 1
         for segment in segments:
             start_time = format_time(float(segment['start']))
@@ -43,9 +44,7 @@ def create_srt_file(segments, output_file):
 
 def resample(file: str, sr: int = 16000):
     """
-    # https://github.com/openai/whisper/blob/7858aa9c08d98f75575035ecd6481f462d66ca27/whisper/audio.py#L22
-    Open an audio file and read as mono waveform, resampling as necessary,
-    save the resampled audio
+    Resample the audio file to 16kHz.
 
     Args:
         file (str): The audio file to open
@@ -54,18 +53,30 @@ def resample(file: str, sr: int = 16000):
     Returns:
         resampled_file (str): The resampled audio file
     """
-    try:
-        # This launches a subprocess to decode audio while down-mixing and resampling as necessary.
-        # Requires the ffmpeg CLI and `ffmpeg-python` package to be installed.
-        out, _ = (
-            ffmpeg.input(file, threads=0)
-            .output("-", format="s16le", acodec="pcm_s16le", ac=1, ar=sr)
-            .run(cmd=["ffmpeg", "-nostdin"], capture_stdout=True, capture_stderr=True)
-        )
-    except ffmpeg.Error as e:
-        raise RuntimeError(f"Failed to load audio: {e.stderr.decode()}") from e
-    np_buffer = np.frombuffer(out, dtype=np.int16)
+    container = av.open(file)
+    stream = next(s for s in container.streams if s.type == 'audio')
 
-    resampled_file = f"{file.split('.')[0]}_resampled.wav"
-    scipy.io.wavfile.write(resampled_file, sr, np_buffer.astype(np.int16))
+    resampler = av.AudioResampler(
+        format='s16',
+        layout='mono',
+        rate=sr,
+    )
+
+    resampled_file = Path(file).stem + "_resampled.wav"
+    output_container = av.open(resampled_file, mode='w')
+    output_stream = output_container.add_stream('pcm_s16le', rate=sr)
+    output_stream.layout = 'mono'
+
+    for frame in container.decode(audio=0):
+        frame.pts = None
+        resampled_frames = resampler.resample(frame)
+        if resampled_frames is not None:
+            for resampled_frame in resampled_frames:
+                for packet in output_stream.encode(resampled_frame):
+                    output_container.mux(packet)
+
+    for packet in output_stream.encode(None):
+        output_container.mux(packet)
+
+    output_container.close()
     return resampled_file