use std::ops::Range;
use crate::segmentation::{Activity, ChunkShape, Segmenter};
use crate::speaker::diarizer::name_at_enrolled_threshold;
use crate::speaker::{AnonymousNumbering, Embedder, Label, Profile};
use super::grid::FrameGrid;
use super::timeline::Timeline;
use super::{Merged, Revision, Spoken, Voices};
use crate::audio::RecordingReader;
const HOP_SAMPLES: u64 = 32_000;
const MIN_TURN: f64 = 0.3;
const MIN_GAP: f64 = 0.5;
const MIN_ACTIVE_FRAMES: usize = 30;
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct Grouping {
pub threshold: f32,
pub least_voiceprints: usize,
}
impl Default for Grouping {
fn default() -> Self {
Self {
threshold: 0.25,
least_voiceprints: 12,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Standing {
pub label: Label,
pub voiceprints: usize,
pub eligible: bool,
}
pub struct PyannoteVoices<S, E> {
segmenter: S,
embedder: E,
enrolled: Vec<Profile>,
grouping: Grouping,
shape: ChunkShape,
grid: FrameGrid,
voices: Vec<Voice>,
next_chunk: u64,
timeline: Timeline,
numbering: AnonymousNumbering,
}
pub struct Chunk {
pub at: u64,
pub frames: Vec<Activity>,
pub voices: Vec<Voiceprint>,
}
pub struct Voiceprint {
pub local: usize,
pub values: Vec<f32>,
}
struct Voice {
label: Label,
sum: Vec<f32>,
count: usize,
}
impl Voice {
fn enrolled(&self) -> bool {
matches!(self.label, Label::Enrolled(_))
}
}
impl<S: Segmenter, E: Embedder> PyannoteVoices<S, E> {
pub fn new(segmenter: S, embedder: E, enrolled: Vec<Profile>, grouping: Grouping) -> Self {
let shape = segmenter.shape();
Self {
segmenter,
embedder,
enrolled,
grouping,
grid: FrameGrid::from(&shape),
shape,
voices: Vec::new(),
next_chunk: 0,
timeline: Timeline::default(),
numbering: AnonymousNumbering::default(),
}
}
pub fn shape(&self) -> &ChunkShape {
&self.shape
}
pub fn hop(&self) -> u64 {
HOP_SAMPLES
}
fn has_chunk(&self, recording: &RecordingReader) -> bool {
match recording.reached_end() {
true => self.next_chunk < recording.held_to(),
false => recording.held_to() >= self.next_chunk + self.shape.samples as u64,
}
}
fn chunk_audio(&self, recording: &RecordingReader, chunk: Range<u64>) -> Option<Vec<f32>> {
let from = chunk.start;
let audio = recording.span(chunk);
if audio.len() == self.shape.samples {
return Some(audio.to_vec());
}
if !recording.reached_end() {
tracing::warn!(
wanted = self.shape.samples,
got = audio.len(),
lost_samples = self.shape.samples - audio.len(),
from,
"audio was released before the voice track read it; \
this stretch will have no speaker"
);
return None;
}
let mut padded = audio.to_vec();
padded.resize(self.shape.samples, 0.0);
Some(padded)
}
pub fn hear(&mut self, recording: &RecordingReader) -> Option<Chunk> {
while self.has_chunk(recording) {
let at = self.next_chunk;
let span = at..at + self.shape.samples as u64;
self.next_chunk += HOP_SAMPLES;
let Some(audio) = self.chunk_audio(recording, span) else {
continue;
};
let frames = match self.segmenter.segment(&audio) {
Ok(frames) => frames,
Err(error) => {
tracing::warn!(
error = %crate::primitives::error_chain(&error),
"speech segmentation failed"
);
continue;
}
};
let voices = (0..self.shape.speakers)
.filter(|&local| {
frames.iter().filter(|frame| frame.speaks(local)).count() >= MIN_ACTIVE_FRAMES
})
.filter_map(|local| {
let heard = masked(&audio, &frames, local, &self.grid);
match self.embedder.embed(&heard) {
Ok(values) => values.map(|values| Voiceprint { local, values }),
Err(error) => {
tracing::warn!(
error = %crate::primitives::error_chain(&error),
"speaker embedding failed"
);
None
}
}
})
.collect();
return Some(Chunk { at, frames, voices });
}
None
}
pub fn group(&mut self, chunk: &Chunk) -> Revision {
let at = self.grid.at(chunk.at);
self.timeline.judge(at, &chunk.frames);
let mut concurrent: Vec<Label> = Vec::new();
for voiceprint in &chunk.voices {
let placed = self.place(&voiceprint.values, &concurrent);
self.timeline
.vote(placed, at, &chunk.frames, voiceprint.local);
concurrent.push(self.voices[placed].label.clone());
}
self.merge(&concurrent)
}
fn place(&mut self, voiceprint: &[f32], concurrent: &[Label]) -> usize {
if let Some(name) = name_at_enrolled_threshold(&self.enrolled, voiceprint) {
let label = Label::Enrolled(name.clone());
return match self.voices.iter().position(|voice| voice.label == label) {
Some(index) => self.absorb(index, voiceprint),
None => self.open(label, voiceprint),
};
}
let nearest = self
.voices
.iter()
.enumerate()
.filter(|(_, voice)| !concurrent.contains(&voice.label))
.map(|(index, voice)| (index, alike(&voice.sum, voice.count, voiceprint, 1)))
.reduce(|best, candidate| {
if candidate.1 > best.1 {
candidate
} else {
best
}
});
if let Some((index, similarity)) = nearest
&& similarity >= self.grouping.threshold
{
return self.absorb(index, voiceprint);
}
let label = self.numbering.issue();
self.open(label, voiceprint)
}
fn absorb(&mut self, index: usize, voiceprint: &[f32]) -> usize {
let voice = &mut self.voices[index];
for (sum, value) in voice.sum.iter_mut().zip(voiceprint) {
*sum += value;
}
voice.count += 1;
index
}
fn open(&mut self, label: Label, voiceprint: &[f32]) -> usize {
self.voices.push(Voice {
label,
sum: voiceprint.to_vec(),
count: 1,
});
self.timeline.open_voice();
self.voices.len() - 1
}
fn merge(&mut self, concurrent: &[Label]) -> Revision {
let mut revision = Revision::default();
while let Some((keeper, absorbed)) = self.closest_pair(concurrent) {
self.timeline.absorb(keeper, absorbed);
let taken = self.voices.remove(absorbed);
revision.merges.push(Merged {
from: taken.label.clone(),
into: self.voices[keeper].label.clone(),
});
let voice = &mut self.voices[keeper];
for (sum, value) in voice.sum.iter_mut().zip(&taken.sum) {
*sum += value;
}
voice.count += taken.count;
}
revision
}
fn closest_pair(&self, concurrent: &[Label]) -> Option<(usize, usize)> {
let mut best: Option<(usize, usize, f32)> = None;
for a in 0..self.voices.len() {
for b in a + 1..self.voices.len() {
if self.apart(a, b, concurrent) {
continue;
}
let similarity = alike(
&self.voices[a].sum,
self.voices[a].count,
&self.voices[b].sum,
self.voices[b].count,
);
if best.is_none_or(|(_, _, top)| similarity > top) {
best = Some((a, b, similarity));
}
}
}
best.filter(|(_, _, similarity)| *similarity >= self.grouping.threshold)
.map(|(a, b, _)| (a, b))
}
fn apart(&self, a: usize, b: usize, concurrent: &[Label]) -> bool {
let (a, b) = (&self.voices[a], &self.voices[b]);
let overlapped = concurrent.contains(&a.label) && concurrent.contains(&b.label);
overlapped || (a.enrolled() && b.enrolled())
}
pub fn heard_so_far(&self) -> Vec<Standing> {
self.voices
.iter()
.map(|voice| Standing {
label: voice.label.clone(),
voiceprints: voice.count,
eligible: voice.count >= self.grouping.least_voiceprints,
})
.collect()
}
fn eligible(&self) -> Vec<bool> {
self.voices
.iter()
.map(|voice| voice.count >= self.grouping.least_voiceprints)
.collect()
}
}
impl<S: Segmenter, E: Embedder> PyannoteVoices<S, E> {
pub fn retained_from(&self) -> u64 {
self.next_chunk
}
pub fn skip_to(&mut self, position: u64) {
self.next_chunk = self.next_chunk.max(position);
}
pub fn turns(&self) -> Vec<Spoken> {
let eligible = self.eligible();
let gap = self.grid.frames(MIN_GAP);
let least = self.grid.frames(MIN_TURN);
let mut spoken: Vec<Spoken> = self
.timeline
.runs(&eligible, gap, least)
.into_iter()
.map(|(voice, run)| Spoken {
span: self.grid.start(run.start)..self.grid.start(run.end),
speaker: self.voices[voice].label.clone(),
})
.collect();
spoken.sort_by_key(|turn| turn.span.start);
spoken
}
}
impl<S: Segmenter, E: Embedder> Voices for PyannoteVoices<S, E> {
fn attribute(&self, spans: &[Range<u64>]) -> Vec<Option<Label>> {
let eligible = self.eligible();
spans
.iter()
.map(|span| {
let opens = self.grid.at(span.start);
let mut owed = vec![0_usize; self.voices.len()];
for frame in opens..self.grid.at(span.end).max(opens + 1) {
for voice in self.timeline.speakers_in(frame, &eligible) {
owed[voice] += 1;
}
}
owed.into_iter()
.enumerate()
.filter(|(_, held)| *held > 0)
.max_by_key(|(_, held)| *held)
.map(|(voice, _)| self.voices[voice].label.clone())
})
.collect()
}
}
fn masked(chunk: &[f32], frames: &[Activity], local: usize, grid: &FrameGrid) -> Vec<f32> {
frames
.iter()
.enumerate()
.filter(|(_, frame)| frame.speaks(local))
.flat_map(|(frame, _)| {
let heard = grid.window(frame);
let from = heard.start.min(chunk.len());
let to = heard.end.min(chunk.len());
chunk[from..to].iter().copied()
})
.collect()
}
fn alike(a: &[f32], a_count: usize, b: &[f32], b_count: usize) -> f32 {
let dot: f32 = a.iter().zip(b).map(|(x, y)| x * y).sum();
dot / (a_count * b_count) as f32
}
#[cfg(test)]
mod testing {
use crate::segmentation::SegmentationError;
use crate::speaker::SpeakerError;
use super::*;
pub(super) struct ScriptedSegmenter {
chunks: std::vec::IntoIter<Vec<Activity>>,
samples: usize,
}
pub(super) const SCRIPTED_SHIFT: usize = 320;
impl Segmenter for ScriptedSegmenter {
fn shape(&self) -> ChunkShape {
ChunkShape {
samples: self.samples,
frame_shift: SCRIPTED_SHIFT,
frame_size: SCRIPTED_SHIFT,
speakers: 3,
}
}
fn segment(&mut self, _chunk: &[f32]) -> Result<Vec<Activity>, SegmentationError> {
Ok(self.chunks.next().unwrap_or_default())
}
}
pub(super) struct ScriptedEmbedder {
voiceprints: std::vec::IntoIter<[f32; 2]>,
}
impl Embedder for ScriptedEmbedder {
fn embed(&mut self, _window: &[f32]) -> Result<Option<Vec<f32>>, SpeakerError> {
Ok(self.voiceprints.next().map(Vec::from))
}
fn dimension(&self) -> usize {
2
}
}
pub(super) fn track(
chunks: Vec<Vec<Activity>>,
voiceprints: Vec<[f32; 2]>,
) -> PyannoteVoices<ScriptedSegmenter, ScriptedEmbedder> {
track_of(HOP_SAMPLES as usize, chunks, voiceprints)
}
pub(super) fn track_with(
grouping: Grouping,
chunks: Vec<Vec<Activity>>,
voiceprints: Vec<[f32; 2]>,
) -> PyannoteVoices<ScriptedSegmenter, ScriptedEmbedder> {
let mut voices = track_of(HOP_SAMPLES as usize, chunks, voiceprints);
voices.grouping = grouping;
voices
}
pub(super) fn track_of(
samples: usize,
chunks: Vec<Vec<Activity>>,
voiceprints: Vec<[f32; 2]>,
) -> PyannoteVoices<ScriptedSegmenter, ScriptedEmbedder> {
PyannoteVoices::new(
ScriptedSegmenter {
chunks: chunks.into_iter(),
samples,
},
ScriptedEmbedder {
voiceprints: voiceprints.into_iter(),
},
Vec::new(),
Grouping::default(),
)
}
pub(super) fn min_voiceprints() -> usize {
Grouping::default().least_voiceprints
}
pub(super) fn hear_all<S: Segmenter, E: Embedder>(
voices: &mut PyannoteVoices<S, E>,
recording: &RecordingReader,
) -> Revision {
let mut revision = Revision::default();
while let Some(chunk) = voices.hear(recording) {
revision.merges.append(&mut voices.group(&chunk).merges);
}
revision
}
pub(super) fn alone(local: usize) -> Vec<Activity> {
vec![Activity::SILENT.speaking(local); MIN_ACTIVE_FRAMES * 2]
}
pub(super) fn recording_of(chunks: usize) -> RecordingReader {
let mut recording =
RecordingReader::from_samples(&vec![0.1; chunks * HOP_SAMPLES as usize]);
recording.fill().expect("the pushed recording is whole");
recording
}
}
#[cfg(test)]
mod tests {
use super::testing::*;
use super::*;
#[test]
fn overlapping_chunks_judge_the_same_frame_more_than_once() {
let width = HOP_SAMPLES as usize * 2;
let per_chunk = width / SCRIPTED_SHIFT;
let heard = min_voiceprints() + 1;
let speaking = vec![Activity::SILENT.speaking(0); per_chunk];
let mut voices = track_of(width, vec![speaking; heard], vec![[1.0, 0.0]; heard]);
hear_all(&mut voices, &recording_of(heard));
let overlap = per_chunk / 2;
let named = voices.timeline.speakers_in(overlap, &voices.eligible());
assert_eq!(named.len(), 1, "the overlap named more than one voice");
}
#[test]
fn a_track_told_to_skip_ahead_never_asks_for_the_audio_behind_it() {
let recording = recording_of(4);
let mut voices = track(
Vec::from([alone(0), alone(0)]),
Vec::from([[1.0, 0.0], [1.0, 0.0]]),
);
voices.skip_to(HOP_SAMPLES * 2);
hear_all(&mut voices, &recording);
assert_eq!(
voices.retained_from(),
HOP_SAMPLES * 4,
"the walk should resume past the skip, not replay it"
);
}
#[test]
fn skipping_backwards_does_nothing() {
let mut voices = track(Vec::new(), Vec::new());
voices.skip_to(HOP_SAMPLES * 3);
voices.skip_to(HOP_SAMPLES);
assert_eq!(voices.retained_from(), HOP_SAMPLES * 3);
}
#[test]
fn two_voiceprints_alike_enough_are_one_voice() {
let recording = recording_of(2);
let mut voices = track(
Vec::from([alone(0), alone(0)]),
Vec::from([[1.0, 0.0], [1.0, 0.0]]),
);
assert!(!hear_all(&mut voices, &recording).revised());
assert_eq!(voices.voices.len(), 1);
}
#[test]
fn a_voiceprint_matching_nothing_opens_a_voice_and_revises_nothing() {
let recording = recording_of(2);
let mut voices = track(
Vec::from([alone(0), alone(0)]),
Vec::from([[1.0, 0.0], [0.0, 1.0]]),
);
assert!(!hear_all(&mut voices, &recording).revised());
assert_eq!(voices.voices.len(), 2);
assert_eq!(voices.voices[0].label, Label::Anonymous(1));
assert_eq!(voices.voices[1].label, Label::Anonymous(2));
}
#[test]
fn an_enrolled_voice_opens_its_own_voice_beside_an_anonymous_one() {
let recording = recording_of(2);
let mut voices = track(
Vec::from([alone(0), alone(0)]),
Vec::from([[1.0, 0.0], [0.0, 1.0]]),
);
let alex = crate::speaker::enrolled_voice(
"alex",
crate::speaker::Role::Guest,
Vec::from([0.0, 1.0]),
);
voices.enrolled = Vec::from([alex.clone()]);
hear_all(&mut voices, &recording);
let live: Vec<&Label> = voices.voices.iter().map(|voice| &voice.label).collect();
assert_eq!(live, [&Label::Anonymous(1), &Label::Enrolled(alex.name)]);
}
#[test]
fn voices_that_drift_together_merge_and_the_earlier_name_survives() {
let recording = recording_of(3);
let mut voices = track(
Vec::from([alone(0), alone(0), alone(0)]),
Vec::from([[1.0, 0.0], [0.17365, 0.98481], [0.5, 0.86603]]),
);
let revision = hear_all(&mut voices, &recording);
assert!(revision.revised());
assert_eq!(revision.merges[0].from, Label::Anonymous(2));
assert_eq!(revision.merges[0].into, Label::Anonymous(1));
assert_eq!(voices.voices.len(), 1);
assert_eq!(voices.voices[0].label, Label::Anonymous(1));
}
#[test]
fn a_merged_number_is_retired_rather_than_handed_to_the_next_voice() {
let recording = recording_of(6);
let mut voices = track(
Vec::from([alone(0), alone(0), alone(0), alone(0), alone(0), alone(0)]),
Vec::from([
[1.0, 0.0],
[-0.1736, 0.9848],
[-0.766, -0.643],
[0.643, 0.766],
[0.342, 0.940],
[0.9, -0.436],
]),
);
let revision = hear_all(&mut voices, &recording);
assert_eq!(revision.merges[0].from, Label::Anonymous(2));
let live: Vec<&Label> = voices.voices.iter().map(|voice| &voice.label).collect();
assert_eq!(
live,
[
&Label::Anonymous(1),
&Label::Anonymous(3),
&Label::Anonymous(4)
]
);
}
#[test]
fn a_span_takes_the_voice_that_owned_the_most_of_it() {
let half = min_voiceprints();
let recording = recording_of(half * 2);
let mut voiceprints = vec![[1.0, 0.0]; half];
voiceprints.extend(vec![[0.0, 1.0]; half]);
let mut voices = track(vec![alone(0); half * 2], voiceprints);
hear_all(&mut voices, &recording);
let turn = HOP_SAMPLES * half as u64;
let named = voices.attribute(&Vec::from([0..turn, turn..turn * 2]));
assert_eq!(
named,
Vec::from([Some(Label::Anonymous(1)), Some(Label::Anonymous(2))])
);
}
#[test]
fn a_voice_heard_too_few_times_names_nothing() {
let brief = min_voiceprints() - 1;
let recording = recording_of(brief);
let mut voices = track(vec![alone(0); brief], vec![[1.0, 0.0]; brief]);
hear_all(&mut voices, &recording);
assert_eq!(voices.voices.len(), 1);
let whole = 0..HOP_SAMPLES * brief as u64;
let named = voices.attribute(&[whole]);
assert_eq!(named, Vec::from([None]));
assert!(voices.turns().is_empty());
}
#[test]
fn lowering_the_bar_lets_a_rare_voice_name_its_audio() {
let brief = min_voiceprints() - 1;
let recording = recording_of(brief);
let mut voices = track_with(
Grouping {
least_voiceprints: brief,
..Grouping::default()
},
vec![alone(0); brief],
vec![[1.0, 0.0]; brief],
);
hear_all(&mut voices, &recording);
let whole = 0..HOP_SAMPLES * brief as u64;
let named = voices.attribute(&[whole]);
assert_eq!(named, Vec::from([Some(Label::Anonymous(1))]));
}
#[test]
fn a_threshold_nothing_can_reach_gives_every_chunk_its_own_voice() {
let heard = min_voiceprints();
let recording = recording_of(heard);
let mut voices = track_with(
Grouping {
threshold: 1.1,
..Grouping::default()
},
vec![alone(0); heard],
vec![[1.0, 0.0]; heard],
);
hear_all(&mut voices, &recording);
assert_eq!(voices.voices.len(), heard, "no two of them could join");
assert!(
voices.heard_so_far().iter().all(|voice| !voice.eligible),
"and none is heard often enough to be written"
);
}
#[test]
fn the_stretch_after_the_last_whole_chunk_is_still_read() {
let whole = min_voiceprints();
let tail = HOP_SAMPLES as usize / 3;
let mut recording =
RecordingReader::from_samples(&vec![0.1; whole * HOP_SAMPLES as usize + tail]);
recording.fill().expect("the pushed recording is whole");
let mut voices = track(vec![alone(0); whole + 1], vec![[1.0, 0.0]; whole + 1]);
hear_all(&mut voices, &recording);
assert!(
voices.retained_from() >= recording.held_to(),
"the track stopped at {} with {} recorded",
voices.retained_from(),
recording.held_to()
);
}
#[test]
fn a_span_past_the_tracks_reach_stays_unnamed() {
let heard = min_voiceprints();
let recording = recording_of(heard);
let mut voices = track(vec![alone(0); heard], vec![[1.0, 0.0]; heard]);
hear_all(&mut voices, &recording);
let past = HOP_SAMPLES * 100..HOP_SAMPLES * 101;
let named = voices.attribute(std::slice::from_ref(&past));
assert_eq!(named, Vec::from([None]));
}
#[test]
fn a_voice_too_brief_to_embed_is_left_alone() {
let recording = recording_of(1);
let brief = vec![Activity::SILENT.speaking(0); MIN_ACTIVE_FRAMES - 1];
let mut voices = track(Vec::from([brief]), Vec::from([[1.0, 0.0]]));
hear_all(&mut voices, &recording);
assert!(voices.voices.is_empty());
}
fn frames(pattern: &[&[usize]]) -> Vec<Activity> {
pattern
.iter()
.map(|speaking| {
speaking
.iter()
.fold(Activity::SILENT, |active, &voice| active.speaking(voice))
})
.collect()
}
fn grid(shift: usize, size: usize) -> FrameGrid {
FrameGrid::from(&ChunkShape {
samples: 8,
frame_shift: shift,
frame_size: size,
speakers: 3,
})
}
#[test]
fn only_a_voices_own_frames_reach_the_embedder() {
let chunk: Vec<f32> = (0..8).map(|sample| sample as f32).collect();
let heard = masked(&chunk, &frames(&[&[], &[0], &[1], &[0]]), 0, &grid(2, 2));
assert_eq!(heard, Vec::from([2.0, 3.0, 6.0, 7.0]));
}
#[test]
fn a_frame_reaches_past_its_own_shift_and_masks_from_the_middle() {
let chunk: Vec<f32> = (0..8).map(|sample| sample as f32).collect();
let heard = masked(&chunk, &frames(&[&[], &[0]]), 0, &grid(2, 4));
assert_eq!(heard, Vec::from([3.0, 4.0]));
}
#[test]
fn a_frame_nobody_speaks_in_masks_to_nothing() {
let chunk = vec![1.0_f32; 8];
assert!(masked(&chunk, &frames(&[&[], &[]]), 0, &grid(2, 2)).is_empty());
}
#[test]
fn one_voice_is_as_alike_to_itself_as_a_cosine_gets() {
let voice = [0.6, 0.8];
assert!((alike(&voice, 1, &voice, 1) - 1.0).abs() < 1e-6);
}
#[test]
fn averaging_over_members_keeps_a_spread_voice_honest() {
let spread = [1.0, 1.0];
let one = [1.0, 0.0];
assert!((alike(&spread, 2, &one, 1) - 0.5).abs() < 1e-6);
}
}