summaryrefslogtreecommitdiff
path: root/src/mic_clipper/vad.py
blob: e82d15afe8f83c71f7de12756114993e7be3e758 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
"""Stateful, local Silero VAD ONNX inference."""

from __future__ import annotations

from importlib import resources
from pathlib import Path

import numpy as np
import onnxruntime


FRAME_SAMPLES = 512
CONTEXT_SAMPLES = 64


class SileroVad:
    """Run the bundled Silero v6 model on exact 32 ms PCM frames."""

    def __init__(self, model_path: Path | None = None) -> None:
        self.model_path = Path(model_path) if model_path else self._default_model_path()
        options = onnxruntime.SessionOptions()
        options.inter_op_num_threads = 1
        options.intra_op_num_threads = 1
        options.enable_cpu_mem_arena = False
        options.log_severity_level = 4
        self.session = onnxruntime.InferenceSession(
            self.model_path,
            providers=["CPUExecutionProvider"],
            sess_options=options,
        )
        self._context = np.zeros(CONTEXT_SAMPLES, dtype=np.float32)
        self._h = np.zeros((1, 1, 128), dtype=np.float32)
        self._c = np.zeros((1, 1, 128), dtype=np.float32)

    def probability(self, samples: np.ndarray) -> float:
        samples = np.asarray(samples, dtype=np.float32)
        if samples.ndim != 1 or len(samples) != FRAME_SAMPLES:
            raise ValueError("Silero VAD requires exactly 512 mono samples at 16 kHz")
        model_input = np.concatenate((self._context, samples)).reshape(1, -1)
        output, self._h, self._c = self.session.run(
            None,
            {"input": model_input, "h": self._h, "c": self._c},
        )
        self._context = samples[-CONTEXT_SAMPLES:].copy()
        return float(output[0])

    @staticmethod
    def _default_model_path() -> Path:
        bundled = resources.files("mic_clipper.assets").joinpath("silero_vad_v6.onnx")
        if not bundled.is_file():
            raise RuntimeError(
                "packaged Silero VAD model is missing: silero_vad_v6.onnx"
            )
        return Path(bundled)