Owner/src/owner_voice_pet/audio_preprocess.py

from __future__ import annotations

from pathlib import Path
from typing import Any

from .models import AudioFrame, ErrorCode, ProviderError
from .speech_models import denoiser_model_path


class NoopAudioPreprocessor:
    def load(self) -> None:
        return None

    def reset(self) -> None:
        return None

    def process_frame(self, frame: AudioFrame) -> AudioFrame:
        return frame

    def flush(self) -> list[AudioFrame]:
        return []


class SherpaOnnxDenoiserPreprocessor:
    def __init__(
        self,
        models_dir: str | Path,
        *,
        sherpa_module: Any | None = None,
    ) -> None:
        self.models_dir = Path(models_dir)
        self._sherpa = sherpa_module
        self._denoiser: Any | None = None
        self.loaded = False

    def load(self) -> None:
        model = denoiser_model_path(self.models_dir)
        if not model.exists():
            raise ProviderError(
                ErrorCode.NOISE_FILTER_MODEL_MISSING,
                f"sherpa-onnx denoiser model is missing: {model}",
                False,
                "sherpa-onnx-gtcrn",
                "audio-preprocess",
            )
        sherpa_onnx = self._sherpa
        if sherpa_onnx is None:
            try:
                import sherpa_onnx  # type: ignore[import-not-found]
            except Exception as exc:
                raise ProviderError(
                    ErrorCode.NOISE_FILTER_FAILED,
                    f"sherpa_onnx is not available: {exc}",
                    False,
                    "sherpa-onnx-gtcrn",
                    "audio-preprocess",
                ) from exc
        try:
            gtcrn = sherpa_onnx.OfflineSpeechDenoiserGtcrnModelConfig(model=str(model))
            model_config = sherpa_onnx.OfflineSpeechDenoiserModelConfig(
                gtcrn=gtcrn,
                num_threads=1,
                provider="cpu",
            )
            config = sherpa_onnx.OnlineSpeechDenoiserConfig(model=model_config)
            self._denoiser = sherpa_onnx.OnlineSpeechDenoiser(config)
        except Exception as exc:
            raise ProviderError(
                ErrorCode.NOISE_FILTER_FAILED,
                f"failed to load sherpa-onnx denoiser: {exc}",
                False,
                "sherpa-onnx-gtcrn",
                "audio-preprocess",
            ) from exc
        self.loaded = True

    def reset(self) -> None:
        if self._denoiser is not None and hasattr(self._denoiser, "reset"):
            self._denoiser.reset()

    def process_frame(self, frame: AudioFrame) -> AudioFrame:
        if not self.loaded or self._denoiser is None:
            raise ProviderError(
                ErrorCode.NOISE_FILTER_FAILED,
                "sherpa-onnx denoiser is not loaded",
                False,
                "sherpa-onnx-gtcrn",
                "audio-preprocess",
            )
        try:
            import numpy as np

            samples = np.frombuffer(frame.pcm, dtype=np.int16).astype(np.float32) / 32768.0
            if frame.channels > 1 and samples.size:
                samples = samples.reshape(-1, frame.channels).mean(axis=1)
            denoised = self._denoiser.run(samples, frame.sample_rate)
            output_samples = np.asarray(getattr(denoised, "samples"), dtype=np.float32)
            output_sample_rate = int(getattr(denoised, "sample_rate", frame.sample_rate))
            clipped = np.clip(output_samples, -1.0, 1.0)
            pcm = (clipped * 32767.0).astype(np.int16).tobytes()
        except ProviderError:
            raise
        except Exception as exc:
            raise ProviderError(
                ErrorCode.NOISE_FILTER_FAILED,
                f"sherpa-onnx denoiser failed: {exc}",
                True,
                "sherpa-onnx-gtcrn",
                "audio-preprocess",
            ) from exc
        metadata = dict(frame.metadata)
        metadata["denoised"] = True
        metadata["noise_filter_provider"] = "sherpa_onnx_gtcrn"
        return AudioFrame(
            pcm=pcm,
            sample_rate=output_sample_rate,
            channels=1,
            timestamp_ms=frame.timestamp_ms,
            frame_id=frame.frame_id,
            metadata=metadata,
        )

    def flush(self) -> list[AudioFrame]:
        if self._denoiser is None or not hasattr(self._denoiser, "flush"):
            return []
        try:
            import numpy as np

            denoised = self._denoiser.flush()
            samples = np.asarray(getattr(denoised, "samples"), dtype=np.float32)
            if samples.size == 0:
                return []
            sample_rate = int(getattr(denoised, "sample_rate", 16000))
            pcm = (np.clip(samples, -1.0, 1.0) * 32767.0).astype(np.int16).tobytes()
        except Exception as exc:
            raise ProviderError(
                ErrorCode.NOISE_FILTER_FAILED,
                f"sherpa-onnx denoiser flush failed: {exc}",
                True,
                "sherpa-onnx-gtcrn",
                "audio-preprocess",
            ) from exc
        return [
            AudioFrame(
                pcm=pcm,
                sample_rate=sample_rate,
                channels=1,
                timestamp_ms=0,
                frame_id=0,
                metadata={
                    "denoised": True,
                    "noise_filter_provider": "sherpa_onnx_gtcrn",
                    "flush": True,
                },
            )
        ]