aria_kit/segmentation/
mod.rs

1//! Who is speaking *when*: the contract every speech-segmentation engine
2//! implements.
3//!
4//! An embedder answers "whose voice is this window" and needs the window to
5//! hold one person. A segmenter answers the question underneath: given a
6//! fixed stretch of audio, which voices are active in each frame of it, and
7//! where does one stop and the next begin. Two people talking at once is an
8//! ordinary answer here rather than a window nobody can embed.
9//!
10//! What a caller does with those frames — masking audio per voice, grouping
11//! the voices across a recording, naming them — is policy, not an engine's
12//! business.
13
14use std::path::PathBuf;
15
16use crate::backend::{BackendError, ModelDownload, ModelError};
17
18/// Speakers active in each frame of one fixed-width chunk of audio.
19///
20/// Implementations map their own failures through
21/// [`BackendError::engine`](crate::backend::BackendError::engine) and name
22/// themselves with [`Engine::NAME`].
23pub trait Segmenter {
24    /// The chunk width this engine reads and the frame grid it answers on.
25    /// Fixed for the life of the engine — it is a property of the graph.
26    fn shape(&self) -> ChunkShape;
27
28    /// Who speaks in each frame of `chunk`, which must be exactly
29    /// [`ChunkShape::samples`] long. A caller with less audio than that pads
30    /// it; a caller with more cuts it.
31    fn segment(&mut self, chunk: &[f32]) -> Result<Vec<Activity>, SegmentationError>;
32}
33
34/// How an engine is constructed from its own configuration. Implementing
35/// this beside [`Segmenter`] is the whole checklist for a new engine.
36pub trait Engine: Segmenter + Sized {
37    const NAME: &'static str;
38
39    /// The files to download for this engine, all into one folder. Empty
40    /// when nothing is downloaded for it. An engine that reads a second
41    /// folder lists that folder's files in a const of its own.
42    const MODEL: &'static [ModelDownload];
43
44    type Config;
45
46    /// Resolves what a load would read, without reading it.
47    fn check_from(config: &Self::Config) -> Result<PathBuf, SegmentationError>;
48
49    fn load_from(config: Self::Config) -> Result<Self, SegmentationError>;
50}
51
52/// The grid one engine answers on.
53///
54/// `frame_size` is wider than `frame_shift` — frames overlap — so a frame's
55/// time is the centre of the audio behind it rather than its leading edge.
56#[derive(Debug, Clone, Copy, PartialEq, Eq)]
57pub struct ChunkShape {
58    /// Samples in one chunk, at 16 kHz.
59    pub samples: usize,
60    /// Samples from one frame to the next.
61    pub frame_shift: usize,
62    /// Samples one frame answers for.
63    pub frame_size: usize,
64    /// Voices the engine can tell apart inside a single chunk. They are
65    /// local to the chunk: voice 0 of one chunk and voice 0 of the next are
66    /// not the same person.
67    pub speakers: usize,
68}
69
70/// Which of a chunk's local voices speak in one frame.
71///
72/// A bitset rather than a `Vec<bool>`: there is one of these per frame and a
73/// ten-second chunk holds hundreds, so the whole answer stays in cache.
74#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
75pub struct Activity(u8);
76
77impl Activity {
78    /// Nobody speaking.
79    pub const SILENT: Self = Self(0);
80
81    /// The most local voices any engine may report. Beyond this an activity
82    /// would not fit its bitset, and no published segmenter comes close.
83    pub const MAX_SPEAKERS: usize = 8;
84
85    /// Marks `local` as speaking. Positions at or past [`Self::MAX_SPEAKERS`]
86    /// are ignored — an engine claiming more voices than the set holds is a
87    /// bug in that engine, and dropping the overflow keeps the frames it did
88    /// get right.
89    pub fn speaking(mut self, local: usize) -> Self {
90        if local < Self::MAX_SPEAKERS {
91            self.0 |= 1 << local;
92        }
93        self
94    }
95
96    pub fn speaks(self, local: usize) -> bool {
97        local < Self::MAX_SPEAKERS && self.0 & (1 << local) != 0
98    }
99
100    /// How many voices are speaking at once.
101    pub fn voices(self) -> usize {
102        self.0.count_ones() as usize
103    }
104}
105
106/// The shared model validator, answering in this capability's vocabulary.
107#[cfg(not(target_family = "wasm"))]
108pub fn validate_model_file(
109    path: PathBuf,
110    backend: &'static str,
111) -> Result<PathBuf, SegmentationError> {
112    crate::backend::validate_model_file(path)
113        .map_err(|source| SegmentationError::Model { backend, source })
114}
115
116#[derive(Debug, thiserror::Error)]
117pub enum SegmentationError {
118    #[error("the {backend} model cannot be used")]
119    Model {
120        backend: &'static str,
121        #[source]
122        source: ModelError,
123    },
124
125    #[error("failed to load the {backend} model {model}")]
126    Load {
127        backend: &'static str,
128        model: String,
129        #[source]
130        source: BackendError,
131    },
132
133    #[error("{backend} inference failed")]
134    Inference {
135        backend: &'static str,
136        #[source]
137        source: BackendError,
138    },
139}
140
141#[cfg(test)]
142mod tests {
143    use super::*;
144
145    #[test]
146    fn silence_names_nobody() {
147        assert_eq!(Activity::SILENT.voices(), 0);
148        assert!(!Activity::SILENT.speaks(0));
149    }
150
151    #[test]
152    fn two_voices_at_once_is_an_ordinary_answer() {
153        let both = Activity::SILENT.speaking(0).speaking(2);
154        assert_eq!(both.voices(), 2);
155        assert!(both.speaks(0) && both.speaks(2));
156        assert!(!both.speaks(1));
157    }
158
159    #[test]
160    fn a_voice_the_set_cannot_hold_is_dropped_not_folded_onto_another() {
161        let overflowed = Activity::SILENT.speaking(Activity::MAX_SPEAKERS);
162        assert_eq!(overflowed.voices(), 0);
163        assert!(!overflowed.speaks(0));
164    }
165}
aria_kit/segmentation/mod.rs

aria_kit/segmentation/
mod.rs
