backend/gate.py

from __future__ import annotations
import time
import numpy as np
import webrtcvad
import config
from backend.types import StageResult
_silero_model = None
_silero_utils = None

def _get_silero():
    global _silero_model, _silero_utils
    if _silero_model is None:
        import torch
        _silero_model, _silero_utils = torch.hub.load(repo_or_dir='snakers4/silero-vad', model='silero_vad', trust_repo=True)
    return (_silero_model, _silero_utils)

def _float_audio(audio: np.ndarray) -> np.ndarray:
    audio = np.asarray(audio)
    if audio.dtype.kind == 'i':
        return audio.astype(np.float32) / 32768.0
    return audio.astype(np.float32)

def _pcm16_bytes(audio_f32: np.ndarray) -> bytes:
    clipped = np.clip(audio_f32, -1.0, 1.0)
    return (clipped * 32767.0).astype(np.int16).tobytes()

def run_webrtcvad(audio: np.ndarray, sample_rate: int=config.SAMPLE_RATE, aggressiveness: int=2) -> StageResult:
    start = time.perf_counter()
    audio_f32 = _float_audio(audio)
    pcm = _pcm16_bytes(audio_f32)
    frame_ms = 30
    frame_bytes = int(sample_rate * (frame_ms / 1000.0)) * 2
    vad = webrtcvad.Vad(aggressiveness)
    speech_frames = 0
    total_frames = 0
    for offset in range(0, len(pcm) - frame_bytes + 1, frame_bytes):
        frame = pcm[offset:offset + frame_bytes]
        total_frames += 1
        if vad.is_speech(frame, sample_rate):
            speech_frames += 1
    speech_detected = total_frames > 0 and speech_frames / total_frames > 0.1
    timing_ms = (time.perf_counter() - start) * 1000
    return StageResult(stage='gate.webrtcvad', timing_ms=timing_ms, output={'speech_detected': speech_detected, 'speech_frame_ratio': speech_frames / total_frames if total_frames else 0.0}, available=True, provenance='real_checkpoint')

def run_silero(audio: np.ndarray, sample_rate: int=config.SAMPLE_RATE, threshold: float=0.5) -> StageResult:
    import torch
    start = time.perf_counter()
    model, utils = _get_silero()
    get_speech_timestamps = utils[0]
    audio_f32 = _float_audio(audio)
    tensor = torch.from_numpy(audio_f32)
    timestamps = get_speech_timestamps(tensor, model, sampling_rate=sample_rate, threshold=threshold)
    speech_detected = len(timestamps) > 0
    timing_ms = (time.perf_counter() - start) * 1000
    return StageResult(stage='gate.silero_vad', timing_ms=timing_ms, output={'speech_detected': speech_detected, 'speech_segments': timestamps}, available=True, provenance='real_checkpoint')

def run_none(audio: np.ndarray, sample_rate: int=config.SAMPLE_RATE) -> StageResult:
    return StageResult(stage='gate.none', timing_ms=0.0, output={'speech_detected': True}, available=True, provenance='rule')

def run(gate: str, audio: np.ndarray, sample_rate: int=config.SAMPLE_RATE, **kwargs) -> StageResult:
    if gate == 'webrtcvad':
        return run_webrtcvad(audio, sample_rate, aggressiveness=kwargs.get('vad_aggressiveness', 2))
    if gate == 'silero_vad':
        return run_silero(audio, sample_rate)
    if gate == 'none':
        return run_none(audio, sample_rate)
    raise ValueError(f'unknown gate: {gate!r}')