diff --git a/src/pipeline/chunk_embedding.rs b/src/pipeline/chunk_embedding.rs index 17a8555..f721d86 100644 --- a/src/pipeline/chunk_embedding.rs +++ b/src/pipeline/chunk_embedding.rs @@ -408,6 +408,7 @@ pub(super) fn try_batch_chunk_embedding( hard_clusters: ChunkSpeakerClusters(Array2::zeros((0, 0))), discrete_diarization: DiscreteDiarization(Array2::zeros((0, 0))), segments: Vec::new(), + exclusive_segments: Vec::new(), }) }) .collect()) diff --git a/src/pipeline/post_inference.rs b/src/pipeline/post_inference.rs index b869ea5..755f54d 100644 --- a/src/pipeline/post_inference.rs +++ b/src/pipeline/post_inference.rs @@ -3,7 +3,7 @@ use tracing::debug; use crate::binarize::binarize; use crate::clustering::plda::PldaTransform; -use crate::reconstruct::Reconstructor; +use crate::reconstruct::{Reconstructor, exclusive_from, resolve_exclusive_conflicts}; use crate::segment::merge_segments; use super::config::{PipelineConfig, ReconstructMethod}; @@ -36,6 +36,7 @@ pub fn post_inference( hard_clusters: ChunkSpeakerClusters(Array2::zeros((0, 0))), discrete_diarization: DiscreteDiarization(Array2::zeros((0, 0))), segments: Vec::new(), + exclusive_segments: Vec::new(), }); } @@ -44,24 +45,38 @@ pub fn post_inference( let reconstructor = Reconstructor::with_clusters(&segmentations, &hard_clusters, &layout.start_frames, 0); + // One activation pass feeds both reconstructions; the exclusive variant needs the + // continuous scores, which a reconstruction has already flattened to 1.0. + let activations = reconstructor.frame_activations(&speaker_count); let discrete_diarization = match config.reconstruct_method { ReconstructMethod::Smoothed { epsilon } => { - reconstructor.reconstruct_smoothed(&speaker_count, epsilon) + reconstructor.reconstruct_smoothed_with(&activations, &speaker_count, epsilon) + } + ReconstructMethod::Standard => { + reconstructor.reconstruct_with(&activations, &speaker_count) } - ReconstructMethod::Standard => reconstructor.reconstruct(&speaker_count), }; + let exclusive_diarization = exclusive_from(&discrete_diarization, &activations); // apply min-duration filtering to remove single-frame speaker flickers let has_duration_filter = config.binarize.min_duration_on > 0 || config.binarize.min_duration_off > 0; - let discrete_diarization = if has_duration_filter { - DiscreteDiarization(binarize(&discrete_diarization, &config.binarize)) + let (discrete_diarization, exclusive_diarization) = if has_duration_filter { + ( + DiscreteDiarization(binarize(&discrete_diarization, &config.binarize)), + DiscreteDiarization(binarize(&exclusive_diarization, &config.binarize)), + ) } else { - discrete_diarization + (discrete_diarization, exclusive_diarization) }; + // binarize runs per-speaker independently and can pad/extend two speakers' regions into the + // same frame, undoing the exclusivity exclusive_from established — re-resolve any conflicts. + let exclusive_diarization = resolve_exclusive_conflicts(&exclusive_diarization, &activations); let segments = discrete_diarization.to_segments(); let segments = merge_segments(&segments, config.merge_gap); + let exclusive_segments = exclusive_diarization.to_segments(); + let exclusive_segments = merge_segments(&exclusive_segments, config.merge_gap); debug!( post_inference_ms = post_start.elapsed().as_millis(), @@ -75,5 +90,6 @@ pub fn post_inference( hard_clusters, discrete_diarization, segments, + exclusive_segments, }) } diff --git a/src/pipeline/types/data.rs b/src/pipeline/types/data.rs index 6ca4c04..e9e427f 100644 --- a/src/pipeline/types/data.rs +++ b/src/pipeline/types/data.rs @@ -172,6 +172,10 @@ pub struct DiarizationResult { pub discrete_diarization: DiscreteDiarization, /// Merged speaker segments (time-stamped speaker turns) pub segments: Vec, + /// Merged speaker segments with at most one speaker at any instant — the + /// `exclusive_speaker_diarization` equivalent. Overlapped frames go to the speaker with + /// the highest activation score, so the speech is kept rather than split or dropped. + pub exclusive_segments: Vec, } impl DiarizationResult { diff --git a/src/reconstruct.rs b/src/reconstruct.rs index ce401c3..d756dc9 100644 --- a/src/reconstruct.rs +++ b/src/reconstruct.rs @@ -134,29 +134,43 @@ impl<'a> Reconstructor<'a> { FrameActivations(activations) } + /// Convenience wrapper over [`reconstruct_with`](Self::reconstruct_with); production code + /// computes activations once and shares them with the exclusive pass. + #[cfg(test)] pub fn reconstruct(&self, speaker_count: &SpeakerCountTrack) -> DiscreteDiarization { - let activations = self.frame_activations(speaker_count); + self.reconstruct_with(&self.frame_activations(speaker_count), speaker_count) + } + + /// [`reconstruct`](Self::reconstruct) over activations the caller already computed — + /// so a caller that also needs the exclusive variant pays for one activation pass. + pub(crate) fn reconstruct_with( + &self, + activations: &FrameActivations, + speaker_count: &SpeakerCountTrack, + ) -> DiscreteDiarization { let mut discrete = Array2::::zeros(activations.raw_dim()); for (frame_idx, &count) in speaker_count.iter().enumerate() { - for speaker_idx in top_k_indices(&activations, frame_idx, count) { + for speaker_idx in top_k_indices(activations, frame_idx, count) { discrete[[frame_idx, speaker_idx]] = 1.0; } } DiscreteDiarization(discrete) } - pub fn reconstruct_smoothed( + /// Smoothed reconstruction over caller-supplied activations, so a caller that also needs + /// the exclusive variant pays for one activation pass. + pub(crate) fn reconstruct_smoothed_with( &self, + activations: &FrameActivations, speaker_count: &SpeakerCountTrack, epsilon: f32, ) -> DiscreteDiarization { - let activations = self.frame_activations(speaker_count); let mut discrete = Array2::::zeros(activations.raw_dim()); let mut previous_speakers: Vec = Vec::new(); for (frame_idx, &count) in speaker_count.iter().enumerate() { let current_speakers = - top_k_indices_smoothed(&activations, frame_idx, count, &previous_speakers, epsilon); + top_k_indices_smoothed(activations, frame_idx, count, &previous_speakers, epsilon); for &speaker_idx in ¤t_speakers { discrete[[frame_idx, speaker_idx]] = 1.0; } @@ -167,6 +181,81 @@ impl<'a> Reconstructor<'a> { } } +/// Collapse a reconstruction to one speaker per frame, keeping the speaker whose *activation +/// score* is highest among those the reconstruction marked active. +/// +/// This is the `exclusive_speaker_diarization` equivalent, and it has to read the continuous +/// activations to mean anything: a reconstruction stores 1.0 for every active speaker, so an +/// argmax taken over it is a tie that resolves to whichever cluster index the comparator +/// happens to favour — the choice ends up made by cluster numbering rather than by acoustics. +/// +/// Speech is never invented or lost: a frame with at least one active speaker keeps exactly +/// one, and a frame with none stays empty. +pub(crate) fn exclusive_from( + full: &DiscreteDiarization, + activations: &FrameActivations, +) -> DiscreteDiarization { + let mut discrete = Array2::::zeros(full.raw_dim()); + for (frame_idx, row) in full.rows().into_iter().enumerate() { + let mut best: Option<(usize, f32)> = None; + for (speaker_idx, &value) in row.iter().enumerate() { + if value <= 0.0 { + continue; + } + let score = activations + .get([frame_idx, speaker_idx]) + .copied() + .unwrap_or(0.0); + if best.is_none_or(|(_, best_score)| score > best_score) { + best = Some((speaker_idx, score)); + } + } + if let Some((speaker_idx, _)) = best { + discrete[[frame_idx, speaker_idx]] = 1.0; + } + } + DiscreteDiarization(discrete) +} + +/// `exclusive_from` guarantees at most one active speaker per frame, but that guarantee doesn't +/// survive `binarize`: it runs per-speaker independently, so `min_duration_on`/`pad_onset`/etc. +/// can extend one speaker's region into a frame another speaker's region was independently +/// extended into, reintroducing overlap in the "exclusive" output. Re-collapse any such frame to +/// its highest-activation speaker, using the same scores `exclusive_from` used originally. +pub(crate) fn resolve_exclusive_conflicts( + exclusive: &DiscreteDiarization, + activations: &FrameActivations, +) -> DiscreteDiarization { + let mut discrete = exclusive.0.clone(); + let num_speakers = discrete.ncols(); + for frame_idx in 0..discrete.nrows() { + let mut best: Option<(usize, f32)> = None; + let mut active_count = 0usize; + for speaker_idx in 0..num_speakers { + if discrete[[frame_idx, speaker_idx]] > 0.0 { + active_count += 1; + let score = activations + .get([frame_idx, speaker_idx]) + .copied() + .unwrap_or(0.0); + if best.is_none_or(|(_, best_score)| score > best_score) { + best = Some((speaker_idx, score)); + } + } + } + if active_count > 1 + && let Some((winner, _)) = best + { + for speaker_idx in 0..num_speakers { + if speaker_idx != winner { + discrete[[frame_idx, speaker_idx]] = 0.0; + } + } + } + } + DiscreteDiarization(discrete) +} + /// Zero out all but the highest-scoring speaker in each frame, making activations exclusive pub fn make_exclusive(activations: &mut Array2) { for mut row in activations.rows_mut() {