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

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

.github/workflows/ci.yml

@@ -35,7 +35,7 @@ jobs:
             ${{ runner.os }}-pip-${{ matrix.python-version }}-
       
       - name: Install system dependencies
-        run: sudo apt-get update && sudo apt-get install -y ffmpeg portaudio19-dev
+        run: sudo apt-get update && sudo apt-get install -y portaudio19-dev
       
       - name: Install Python dependencies
         run: |
@@ -180,7 +180,7 @@ jobs:
             ubuntu-latest-pip-3.8-
 
       - name: Install system dependencies
-        run: sudo apt-get update && sudo apt-get install -y ffmpeg portaudio19-dev
+        run: sudo apt-get update && sudo apt-get install -y portaudio19-dev
 
       - name: Install Python dependencies
         run: |

README.md

@@ -12,7 +12,7 @@ to convert speech input into text output. It can be used to transcribe both live
 input from microphone and pre-recorded audio files.
 
 ## Installation
-- Install PyAudio and ffmpeg
+- Install PyAudio
 ```bash
  bash scripts/setup.sh
 ```

requirements/client.txt

@@ -1,4 +1,4 @@
 PyAudio
-ffmpeg-python
+av
 scipy
 websocket-client

requirements/server.txt

@@ -4,8 +4,8 @@ onnxruntime==1.16.0
 numba
 kaldialign
 soundfile
-ffmpeg-python
 scipy
+av
 jiwer
 evaluate
 numpy<2

scripts/setup.sh

@@ -1,3 +1,3 @@
 #! /bin/bash
 
-apt-get install portaudio19-dev ffmpeg wget -y
+apt-get install portaudio19-dev wget -y

setup.py

@@ -48,7 +48,6 @@ setup(
         "torchaudio",
         "websockets",
         "onnxruntime==1.16.0",
-        "ffmpeg-python",
         "scipy",
         "websocket-client",
         "numba",

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)
+        print("[INFO]: Connecting to RTSP stream...")
-        self.handle_ffmpeg_process(process, stream_type='RTSP')
+        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)
+        print("[INFO]: Connecting to HLS stream...")
-        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()
         try:
-            # Process the stream
+            container = av.open(hls_url, format="hls")
-            while True:
+            self.process_av_stream(container, stream_type="HLS", save_file=save_file)
-                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())
-
         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:
+        print("[INFO]: HLS stream processing finished.")
-                process.kill()
 
-        print(f"[INFO]: {stream_type} stream processing finished.")
+    def process_av_stream(self, container, stream_type, save_file=None):
-
-    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):
         """
-        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""):
+        audio_stream = next((s for s in container.streams if s.type == "audio"), None)
-            logging.debug(f'[STDERR]: {line.decode()}')
+        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
+    with wave.open(resampled_file, "rb") as wav_file:
-    # and resampling as necessary.  Requires the ffmpeg CLI in PATH.
+        num_frames = wav_file.getnframes()
-    # fmt: off
+        raw_data = wav_file.readframes(num_frames)
-    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
 
-    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):
+def create_srt_file(segments, resampled_file):
-    with open(output_file, 'w', encoding='utf-8') as srt_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
+    Resample the audio file to 16kHz.
-    Open an audio file and read as mono waveform, resampling as necessary,
-    save the resampled audio
 
     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:
+    container = av.open(file)
-        # This launches a subprocess to decode audio while down-mixing and resampling as necessary.
+    stream = next(s for s in container.streams if s.type == 'audio')
-        # 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)
 
-    resampled_file = f"{file.split('.')[0]}_resampled.wav"
+    resampler = av.AudioResampler(
-    scipy.io.wavfile.write(resampled_file, sr, np_buffer.astype(np.int16))
+        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