File size: 1,046 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 | """Postprocessing for DiariZen segmentation output."""
import numpy as np
def log_probs_to_probs(log_probs: np.ndarray) -> np.ndarray:
"""Convert log probabilities to probabilities via softmax.
Args:
log_probs: (1, frames, 11) log-softmax output.
Returns:
(1, frames, 11) probability distribution.
"""
max_val = log_probs.max(axis=-1, keepdims=True)
exp_vals = np.exp(log_probs - max_val)
return exp_vals / exp_vals.sum(axis=-1, keepdims=True)
def top_speakers_at_frame(
log_probs: np.ndarray,
frame_idx: int,
top_k: int = 3,
) -> list[tuple[int, float]]:
"""Get top-k speaker class indices and their log-probabilities at a frame.
Args:
log_probs: (1, frames, 11) output.
frame_idx: Frame index.
top_k: Number of top classes.
Returns:
List of (class_index, log_probability) tuples.
"""
frame = log_probs[0, frame_idx]
top_indices = np.argsort(frame)[-top_k:][::-1]
return [(int(i), float(frame[i])) for i in top_indices]
|