| """DiariZen speaker diarization segmentation inference.""" |
|
|
| import json |
| import os |
| from pathlib import Path |
| from typing import Optional |
|
|
| import numpy as np |
|
|
|
|
| class DiarizenSegmenter: |
| """Speaker diarization segmentation using NPU CNN frontend + CPU backend. |
| |
| Pipeline: |
| 1. Audio preprocessing (resample + LayerNorm) on CPU. |
| 2. CNN feature extraction on AX650 NPU (U16). |
| 3. Transformer + Conformer + Classifier on CPU (ONNX Runtime). |
| """ |
|
|
| def __init__( |
| self, |
| cnn_model_path: str, |
| backend_model_path: str, |
| meta_path: Optional[str] = None, |
| ): |
| self._cnn_path = Path(cnn_model_path) |
| self._backend_path = Path(backend_model_path) |
|
|
| if meta_path is None: |
| meta_path = Path(__file__).parent / "model_meta.json" |
| with open(meta_path) as f: |
| self._meta = json.load(f) |
|
|
| pp = self._meta["preprocess"] |
| self._sample_rate = pp["sample_rate"] |
| self._duration_s = pp["duration_seconds"] |
| self._eps = pp.get("layer_norm_eps", 1e-5) |
| self._num_samples = int(self._sample_rate * self._duration_s) |
|
|
| self._cnn_session = None |
| self._backend_session = None |
|
|
| def _init_cnn(self): |
| """Initialize NPU CNN inference session.""" |
| try: |
| from axengine import InferenceSession |
| except ImportError: |
| raise RuntimeError( |
| "pyaxengine is required for NPU inference. " |
| "Install from: https://github.com/AXERA-TECH/pyaxengine" |
| ) |
| self._cnn_session = InferenceSession(str(self._cnn_path)) |
|
|
| def _init_backend(self): |
| """Initialize CPU backend ONNX inference session.""" |
| import onnxruntime as ort |
| self._backend_session = ort.InferenceSession( |
| str(self._backend_path), |
| providers=["CPUExecutionProvider"], |
| ) |
|
|
| def __call__(self, audio: np.ndarray, sample_rate: int) -> np.ndarray: |
| """Run segmentation inference. |
| |
| Args: |
| audio: 1-D float32 waveform. |
| sample_rate: Original sample rate. |
| |
| Returns: |
| Log-probabilities of shape (1, frames, 11), float32. |
| """ |
| from .preprocess import preprocess_audio |
|
|
| waveform_ln = preprocess_audio( |
| audio, sample_rate, |
| target_sr=self._sample_rate, |
| duration_s=self._duration_s, |
| eps=self._eps, |
| ) |
|
|
| if self._cnn_session is None: |
| self._init_cnn() |
| cnn_outputs = self._cnn_session.run( |
| {self._cnn_session.input_names()[0]: waveform_ln} |
| ) |
| cnn_features = cnn_outputs[0] |
|
|
| if self._backend_session is None: |
| self._init_backend() |
| backend_inputs = { |
| self._backend_session.get_inputs()[0].name: cnn_features |
| } |
| log_probs = self._backend_session.run(None, backend_inputs)[0] |
| return log_probs |
|
|
| @property |
| def num_frames(self) -> int: |
| return 199 |
|
|
| @property |
| def num_classes(self) -> int: |
| return 11 |
|
|
