aria_segmentation/pyannote/
mod.rs

1//! Speech segmentation through the pyannote 3.0 ONNX graph.
2//!
3//! Ten seconds of 16 kHz audio in, one answer per 270 samples out: which of
4//! up to three voices speak in that frame. The graph scores seven mutually
5//! exclusive classes — silence, each voice alone, each pair — so the winning
6//! class names the speakers and there is nothing to threshold.
7//!
8//! No output class represents three simultaneous voices.
9//!
10//! ```no_run
11//! use aria_kit::segmentation::{Engine, Segmenter};
12//! use aria_segmentation::pyannote::{Config, Pyannote};
13//!
14//! # fn main() -> Result<(), aria_kit::segmentation::SegmentationError> {
15//! let mut engine = Pyannote::load_from(Config {
16//!     model: "segmentation-3.0.onnx".into(),
17//! })?;
18//! let frames = engine.segment(&vec![0.0; 160_000])?;
19//! # Ok(())
20//! # }
21//! ```
22
23use std::path::PathBuf;
24
25use aria_kit::backend::BackendError;
26use aria_kit::segmentation::{
27    Activity, ChunkShape, Engine, SegmentationError, Segmenter, validate_model_file,
28};
29use ort::session::Session;
30use ort::value::Tensor;
31
32pub const NAME: &str = "pyannote";
33
34/// The chunk the graph was trained on: ten seconds at 16 kHz.
35const WINDOW_SAMPLES: usize = 160_000;
36/// Samples from one output frame to the next.
37const FRAME_SHIFT: usize = 270;
38/// Samples one output frame answers for. Wider than the shift, so frames
39/// overlap and a frame's time is its centre.
40const FRAME_SIZE: usize = 991;
41/// Voices the graph tells apart inside a single chunk.
42const LOCAL_SPEAKERS: usize = 3;
43const AUDIO_INPUT: &str = "x";
44const SCORES_OUTPUT: &str = "y";
45
46/// Which local voices each of the seven classes names: silence, then each
47/// voice alone, then each pair.
48const POWERSET: [&[usize]; 7] = [&[], &[0], &[1], &[2], &[0, 1], &[0, 2], &[1, 2]];
49
50pub struct Config {
51    pub model: PathBuf,
52}
53
54/// A loaded pyannote graph. Holding one is proof the model runs: loading
55/// segments a silent chunk before returning.
56pub struct Pyannote {
57    session: Session,
58}
59
60impl Engine for Pyannote {
61    const NAME: &'static str = NAME;
62
63    type Config = Config;
64
65    fn check_from(config: &Self::Config) -> Result<PathBuf, SegmentationError> {
66        validate_model_file(config.model.clone(), NAME)
67    }
68
69    fn load_from(config: Self::Config) -> Result<Self, SegmentationError> {
70        let model = Self::check_from(&config)?;
71        let session = Session::builder()
72            .and_then(|mut builder| builder.commit_from_file(&model))
73            .map_err(|source| SegmentationError::Load {
74                backend: NAME,
75                model: model.display().to_string(),
76                source: BackendError::engine(source),
77            })?;
78        let mut engine = Self { session };
79        engine.segment(&vec![0.0; WINDOW_SAMPLES])?;
80        Ok(engine)
81    }
82}
83
84impl Segmenter for Pyannote {
85    fn shape(&self) -> ChunkShape {
86        ChunkShape {
87            samples: WINDOW_SAMPLES,
88            frame_shift: FRAME_SHIFT,
89            frame_size: FRAME_SIZE,
90            speakers: LOCAL_SPEAKERS,
91        }
92    }
93
94    fn segment(&mut self, chunk: &[f32]) -> Result<Vec<Activity>, SegmentationError> {
95        if chunk.len() != WINDOW_SAMPLES {
96            return Err(unusable(format!(
97                "a chunk of {} samples reached the model, which reads {WINDOW_SAMPLES}",
98                chunk.len()
99            )));
100        }
101        let shape = [1_i64, 1, WINDOW_SAMPLES as i64];
102        let input = Tensor::from_array((shape, chunk.to_vec())).map_err(inference)?;
103        let mut outputs = self
104            .session
105            .run(ort::inputs![AUDIO_INPUT => input])
106            .map_err(inference)?;
107        let scores = outputs
108            .remove(SCORES_OUTPUT)
109            .ok_or_else(|| unusable(format!("the model returned no {SCORES_OUTPUT} output")))?;
110        let (shape, values) = scores.try_extract_tensor::<f32>().map_err(inference)?;
111        let classes = *shape.last().unwrap_or(&0) as usize;
112        if classes != POWERSET.len() {
113            return Err(unusable(format!(
114                "the model scores {classes} classes, expected {}",
115                POWERSET.len()
116            )));
117        }
118        Ok(values.chunks_exact(classes).map(loudest).collect())
119    }
120}
121
122/// The voices named by a frame's highest-scoring class.
123fn loudest(frame: &[f32]) -> Activity {
124    let winner = frame
125        .iter()
126        .enumerate()
127        .max_by(|a, b| a.1.total_cmp(b.1))
128        .map_or(0, |(class, _)| class);
129    POWERSET[winner]
130        .iter()
131        .fold(Activity::SILENT, |active, &voice| active.speaking(voice))
132}
133
134fn inference(source: ort::Error) -> SegmentationError {
135    SegmentationError::Inference {
136        backend: NAME,
137        source: BackendError::engine(source),
138    }
139}
140
141fn unusable(reason: String) -> SegmentationError {
142    SegmentationError::Inference {
143        backend: NAME,
144        source: BackendError::Unusable(reason),
145    }
146}
147
148#[cfg(test)]
149mod tests {
150    use super::*;
151
152    #[test]
153    fn the_classes_run_silence_then_singles_then_pairs() {
154        assert_eq!(
155            loudest(&[9.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
156            Activity::SILENT
157        );
158        assert_eq!(
159            loudest(&[0.0, 9.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
160            Activity::SILENT.speaking(0)
161        );
162        assert_eq!(
163            loudest(&[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 9.0]),
164            Activity::SILENT.speaking(1).speaking(2)
165        );
166    }
167
168    #[test]
169    fn every_class_names_at_most_two_voices() {
170        for class in POWERSET {
171            assert!(class.len() <= 2);
172            assert!(class.iter().all(|&voice| voice < LOCAL_SPEAKERS));
173        }
174    }
175}
aria_segmentation/pyannote/mod.rs

aria_segmentation/pyannote/
mod.rs
