| """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 |
|
|
