File size: 2,012 Bytes
9ded405 142a394 9ded405 | 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 | """Turn Nemotron 3 frame probabilities into speaker segments, replicating Omi's published decoding behaviour.
Omi's output differs from NVIDIA's plain 0.5 threshold by a slightly higher threshold and by bridging short
same-speaker pauses: always below `bridge_s`, and up to `bridge_clear_s` when no other speaker is active inside.
"""
from dataclasses import dataclass
import numpy as np
FRAME_S = 0.01
@dataclass(frozen=True)
class DecoderParams:
threshold: float = 0.5
bridge_s: float = 0.0
bridge_clear_s: float = 0.0
def _runs(active: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
edges = np.diff(np.r_[False, active, False].astype(np.int8))
return np.flatnonzero(edges == 1), np.flatnonzero(edges == -1)
def activity(probs: np.ndarray, params: DecoderParams) -> np.ndarray:
"""Boolean (frames, speakers) activity after thresholding and pause bridging."""
raw = probs > params.threshold
active = raw.copy()
bridge, bridge_clear = round(params.bridge_s / FRAME_S), round(params.bridge_clear_s / FRAME_S)
for speaker in range(raw.shape[1]):
others = np.delete(raw, speaker, axis=1).any(axis=1)
starts, ends = _runs(raw[:, speaker])
for gap_start, gap_end in zip(ends[:-1], starts[1:]):
length = gap_end - gap_start
if length < bridge or (length < bridge_clear and not others[gap_start:gap_end].any()):
active[gap_start:gap_end, speaker] = True
return active
def decode(probs: np.ndarray, params: DecoderParams) -> list[dict]:
"""Segments in seconds; speakers that never activate are dropped, so the count stays automatic."""
active = activity(probs, params)
segments = []
for speaker in range(active.shape[1]):
starts, ends = _runs(active[:, speaker])
segments.extend(
{"start": round(float(s * FRAME_S), 2), "end": round(float(e * FRAME_S), 2), "speaker": f"spk_{speaker}"}
for s, e in zip(starts, ends)
)
return segments
|