infinite/
infinite.rs

1/*
2wget https://github.com/thewh1teagle/pyannote-rs/releases/download/v0.1.0/segmentation-3.0.onnx
3wget https://github.com/thewh1teagle/pyannote-rs/releases/download/v0.1.0/wespeaker_en_voxceleb_CAM++.onnx
4wget https://github.com/thewh1teagle/pyannote-rs/releases/download/v0.1.0/6_speakers.wav
5cargo run --example infinite 6_speakers.wav
6*/
7
8use pyannote_rs::{EmbeddingExtractor, EmbeddingManager};
9
10fn process_segment(
11    segment: pyannote_rs::Segment,
12    embedding_extractor: &mut EmbeddingExtractor,
13    embedding_manager: &mut EmbeddingManager,
14    search_threshold: f32,
15) -> Result<(), eyre::Report> {
16    let embedding_result: Vec<f32> = embedding_extractor
17        .compute(&segment.samples)
18        .unwrap()
19        .collect();
20
21    let speaker = embedding_manager
22        .search_speaker(embedding_result.clone(), search_threshold)
23        .ok_or_else(|| embedding_manager.search_speaker(embedding_result, 0.0)) // Ensure always to return speaker
24        .map(|r| r.to_string())
25        .unwrap_or("?".into());
26
27    println!(
28        "start = {:.2}, end = {:.2}, speaker = {}",
29        segment.start, segment.end, speaker
30    );
31
32    Ok(())
33}
34
35fn main() -> Result<(), eyre::Report> {
36    let audio_path = std::env::args().nth(1).expect("Please specify audio file");
37    let search_threshold = 0.5;
38
39    let embedding_model_path = "wespeaker_en_voxceleb_CAM++.onnx";
40    let segmentation_model_path = "segmentation-3.0.onnx";
41
42    let (samples, sample_rate) = pyannote_rs::read_wav(&audio_path)?;
43    let mut embedding_extractor = EmbeddingExtractor::new(embedding_model_path)?;
44    let mut embedding_manager = EmbeddingManager::new(usize::MAX);
45
46    let segments = pyannote_rs::get_segments(&samples, sample_rate, segmentation_model_path)?;
47
48    for segment in segments {
49        if let Ok(segment) = segment {
50            if let Err(error) = process_segment(
51                segment,
52                &mut embedding_extractor,
53                &mut embedding_manager,
54                search_threshold,
55            ) {
56                eprintln!("Error processing segment: {:?}", error);
57            }
58        } else if let Err(error) = segment {
59            eprintln!("Failed to process segment: {:?}", error);
60        }
61    }
62
63    Ok(())
64}