File size: 3,088 Bytes
37c4768 | 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 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | """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
|