diff --git a/docker/Dockerfile.riva b/docker/Dockerfile.riva index b153017..63f32c2 100644 --- a/docker/Dockerfile.riva +++ b/docker/Dockerfile.riva @@ -5,7 +5,7 @@ RUN apt-get update && \ apt-get clean && \ rm -rf /var/lib/apt/lists/* -RUN pip install pyaudio==0.2.14 websockets==16.0 numpy==1.26.4 opencv-python==4.11.0.86 requests==2.32.5 google==3.0.0 eclipse-zenoh==1.7.2 pycdr2==1.0.0 av==16.1.0 vulture==2.14 openai==2.24.0 protobuf==6.33.5 +RUN pip install pyaudio==0.2.14 websockets==16.0 numpy==1.26.4 opencv-python==4.11.0.86 requests==2.32.5 google==3.0.0 eclipse-zenoh==1.7.2 pycdr2==1.0.0 av==16.1.0 vulture==2.14 openai==2.24.0 protobuf==6.33.5 prometheus_client RUN pip uninstall -y wandb && pip install wandb==0.18.7 RUN mkdir /app diff --git a/src/om1_speech/riva/audio_stream_input.py b/src/om1_speech/riva/audio_stream_input.py index 6f33e1a..5bf184e 100644 --- a/src/om1_speech/riva/audio_stream_input.py +++ b/src/om1_speech/riva/audio_stream_input.py @@ -1,6 +1,7 @@ import base64 import json import logging +import struct from queue import Empty, Queue from typing import Any, Dict, Optional @@ -40,12 +41,9 @@ def handle_ws_incoming_message(self, connection_id: str, message: Any): expected to be binary audio data """ try: - # Verify we received binary data if isinstance(message, bytes): - logging.error("Legacy audio stream input. Set rate to 1600.") - self.audio_queue.put( - {"audio": base64.b64encode(message).decode("utf-8"), "rate": 16000} - ) + audio, rate = self._parse_binary(message) + self.audio_queue.put({"audio": audio, "rate": rate}) if isinstance(message, str): try: message = json.loads(message) @@ -100,3 +98,21 @@ def stop(self): Sets the running flag to False to stop processing. """ self.running = False + + @staticmethod + def _parse_binary(data: bytes): + """Parse binary format of ASR. + + Falls back to raw PCM at 16 kHz if the header is invalid. + """ + if len(data) > 4: + header_len = struct.unpack(">I", data[:4])[0] + if 4 + header_len < len(data): + try: + header = json.loads(data[4 : 4 + header_len]) + pcm = data[4 + header_len :] + rate = header.get("rate", 16000) + return base64.b64encode(pcm).decode("utf-8"), rate + except (json.JSONDecodeError, UnicodeDecodeError): + pass + return base64.b64encode(data).decode("utf-8"), 16000