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, ModelDownload, Unpack};
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/// Segmentation 3.0 as sherpa-onnx exports it. This engine relies on that
35/// export's input and output names, and only sherpa-onnx's release tarball
36/// ships it.
37pub const MODEL: &[ModelDownload] = &[ModelDownload {
38    file_name: "segmentation-3.0.onnx",
39    url: "https://github.com/k2-fsa/sherpa-onnx/releases/download/speaker-segmentation-models/sherpa-onnx-pyannote-segmentation-3-0.tar.bz2",
40    bytes: 6_958_444,
41    sha256: "24615ee884c897d9d2ba09bb4d30da6bb1b15e685065962db5b02e76e4996488",
42    unpack: Unpack::TarBz2Member("sherpa-onnx-pyannote-segmentation-3-0/model.onnx"),
43}];
44
45/// The chunk the graph was trained on: ten seconds at 16 kHz.
46const WINDOW_SAMPLES: usize = 160_000;
47/// Samples from one output frame to the next.
48const FRAME_SHIFT: usize = 270;
49/// Samples one output frame answers for. Wider than the shift, so frames
50/// overlap and a frame's time is its centre.
51const FRAME_SIZE: usize = 991;
52/// Voices the graph tells apart inside a single chunk.
53const LOCAL_SPEAKERS: usize = 3;
54const AUDIO_INPUT: &str = "x";
55const SCORES_OUTPUT: &str = "y";
56
57/// Which local voices each of the seven classes names: silence, then each
58/// voice alone, then each pair.
59const POWERSET: [&[usize]; 7] = [&[], &[0], &[1], &[2], &[0, 1], &[0, 2], &[1, 2]];
60
61pub struct Config {
62    pub model: PathBuf,
63}
64
65/// A loaded pyannote graph. Holding one is proof the model runs: loading
66/// segments a silent chunk before returning.
67pub struct Pyannote {
68    session: Session,
69}
70
71impl Engine for Pyannote {
72    const NAME: &'static str = NAME;
73    const MODEL: &'static [ModelDownload] = MODEL;
74
75    type Config = Config;
76
77    fn check_from(config: &Self::Config) -> Result<PathBuf, SegmentationError> {
78        validate_model_file(config.model.clone(), NAME)
79    }
80
81    fn load_from(config: Self::Config) -> Result<Self, SegmentationError> {
82        let model = Self::check_from(&config)?;
83        let session = Session::builder()
84            .and_then(|mut builder| builder.commit_from_file(&model))
85            .map_err(|source| SegmentationError::Load {
86                backend: NAME,
87                model: model.display().to_string(),
88                source: BackendError::engine(source),
89            })?;
90        let mut engine = Self { session };
91        engine.segment(&vec![0.0; WINDOW_SAMPLES])?;
92        Ok(engine)
93    }
94}
95
96impl Segmenter for Pyannote {
97    fn shape(&self) -> ChunkShape {
98        ChunkShape {
99            samples: WINDOW_SAMPLES,
100            frame_shift: FRAME_SHIFT,
101            frame_size: FRAME_SIZE,
102            speakers: LOCAL_SPEAKERS,
103        }
104    }
105
106    fn segment(&mut self, chunk: &[f32]) -> Result<Vec<Activity>, SegmentationError> {
107        if chunk.len() != WINDOW_SAMPLES {
108            return Err(unusable(format!(
109                "a chunk of {} samples reached the model, which reads {WINDOW_SAMPLES}",
110                chunk.len()
111            )));
112        }
113        let shape = [1_i64, 1, WINDOW_SAMPLES as i64];
114        let input = Tensor::from_array((shape, chunk.to_vec())).map_err(inference)?;
115        let mut outputs = self
116            .session
117            .run(ort::inputs![AUDIO_INPUT => input])
118            .map_err(inference)?;
119        let scores = outputs
120            .remove(SCORES_OUTPUT)
121            .ok_or_else(|| unusable(format!("the model returned no {SCORES_OUTPUT} output")))?;
122        let (shape, values) = scores.try_extract_tensor::<f32>().map_err(inference)?;
123        let classes = *shape.last().unwrap_or(&0) as usize;
124        if classes != POWERSET.len() {
125            return Err(unusable(format!(
126                "the model scores {classes} classes, expected {}",
127                POWERSET.len()
128            )));
129        }
130        Ok(values.chunks_exact(classes).map(loudest).collect())
131    }
132}
133
134/// The voices named by a frame's highest-scoring class.
135fn loudest(frame: &[f32]) -> Activity {
136    let winner = frame
137        .iter()
138        .enumerate()
139        .max_by(|a, b| a.1.total_cmp(b.1))
140        .map_or(0, |(class, _)| class);
141    POWERSET[winner]
142        .iter()
143        .fold(Activity::SILENT, |active, &voice| active.speaking(voice))
144}
145
146fn inference(source: ort::Error) -> SegmentationError {
147    SegmentationError::Inference {
148        backend: NAME,
149        source: BackendError::engine(source),
150    }
151}
152
153fn unusable(reason: String) -> SegmentationError {
154    SegmentationError::Inference {
155        backend: NAME,
156        source: BackendError::Unusable(reason),
157    }
158}
159
160#[cfg(test)]
161mod tests {
162    use super::*;
163
164    #[test]
165    fn the_classes_run_silence_then_singles_then_pairs() {
166        assert_eq!(
167            loudest(&[9.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
168            Activity::SILENT
169        );
170        assert_eq!(
171            loudest(&[0.0, 9.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
172            Activity::SILENT.speaking(0)
173        );
174        assert_eq!(
175            loudest(&[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 9.0]),
176            Activity::SILENT.speaking(1).speaking(2)
177        );
178    }
179
180    #[test]
181    fn every_class_names_at_most_two_voices() {
182        for class in POWERSET {
183            assert!(class.len() <= 2);
184            assert!(class.iter().all(|&voice| voice < LOCAL_SPEAKERS));
185        }
186    }
187}
aria_segmentation/pyannote/mod.rs

aria_segmentation/pyannote/
mod.rs
