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, 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    type Config;
40
41    /// Resolves what a load would read, without reading it.
42    fn check_from(config: &Self::Config) -> Result<PathBuf, SegmentationError>;
43
44    fn load_from(config: Self::Config) -> Result<Self, SegmentationError>;
45}
46
47/// The grid one engine answers on.
48///
49/// `frame_size` is wider than `frame_shift` — frames overlap — so a frame's
50/// time is the centre of the audio behind it rather than its leading edge.
51#[derive(Debug, Clone, Copy, PartialEq, Eq)]
52pub struct ChunkShape {
53    /// Samples in one chunk, at 16 kHz.
54    pub samples: usize,
55    /// Samples from one frame to the next.
56    pub frame_shift: usize,
57    /// Samples one frame answers for.
58    pub frame_size: usize,
59    /// Voices the engine can tell apart inside a single chunk. They are
60    /// local to the chunk: voice 0 of one chunk and voice 0 of the next are
61    /// not the same person.
62    pub speakers: usize,
63}
64
65/// Which of a chunk's local voices speak in one frame.
66///
67/// A bitset rather than a `Vec<bool>`: there is one of these per frame and a
68/// ten-second chunk holds hundreds, so the whole answer stays in cache.
69#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
70pub struct Activity(u8);
71
72impl Activity {
73    /// Nobody speaking.
74    pub const SILENT: Self = Self(0);
75
76    /// The most local voices any engine may report. Beyond this an activity
77    /// would not fit its bitset, and no published segmenter comes close.
78    pub const MAX_SPEAKERS: usize = 8;
79
80    /// Marks `local` as speaking. Positions at or past [`Self::MAX_SPEAKERS`]
81    /// are ignored — an engine claiming more voices than the set holds is a
82    /// bug in that engine, and dropping the overflow keeps the frames it did
83    /// get right.
84    pub fn speaking(mut self, local: usize) -> Self {
85        if local < Self::MAX_SPEAKERS {
86            self.0 |= 1 << local;
87        }
88        self
89    }
90
91    pub fn speaks(self, local: usize) -> bool {
92        local < Self::MAX_SPEAKERS && self.0 & (1 << local) != 0
93    }
94
95    /// How many voices are speaking at once.
96    pub fn voices(self) -> usize {
97        self.0.count_ones() as usize
98    }
99}
100
101/// The shared model validator, answering in this capability's vocabulary.
102#[cfg(not(target_family = "wasm"))]
103pub fn validate_model_file(
104    path: PathBuf,
105    backend: &'static str,
106) -> Result<PathBuf, SegmentationError> {
107    crate::backend::validate_model_file(path)
108        .map_err(|source| SegmentationError::Model { backend, source })
109}
110
111#[derive(Debug, thiserror::Error)]
112pub enum SegmentationError {
113    #[error("the {backend} model cannot be used")]
114    Model {
115        backend: &'static str,
116        #[source]
117        source: ModelError,
118    },
119
120    #[error("failed to load the {backend} model {model}")]
121    Load {
122        backend: &'static str,
123        model: String,
124        #[source]
125        source: BackendError,
126    },
127
128    #[error("{backend} inference failed")]
129    Inference {
130        backend: &'static str,
131        #[source]
132        source: BackendError,
133    },
134}
135
136#[cfg(test)]
137mod tests {
138    use super::*;
139
140    #[test]
141    fn silence_names_nobody() {
142        assert_eq!(Activity::SILENT.voices(), 0);
143        assert!(!Activity::SILENT.speaks(0));
144    }
145
146    #[test]
147    fn two_voices_at_once_is_an_ordinary_answer() {
148        let both = Activity::SILENT.speaking(0).speaking(2);
149        assert_eq!(both.voices(), 2);
150        assert!(both.speaks(0) && both.speaks(2));
151        assert!(!both.speaks(1));
152    }
153
154    #[test]
155    fn a_voice_the_set_cannot_hold_is_dropped_not_folded_onto_another() {
156        let overflowed = Activity::SILENT.speaking(Activity::MAX_SPEAKERS);
157        assert_eq!(overflowed.voices(), 0);
158        assert!(!overflowed.speaks(0));
159    }
160}
aria_kit/segmentation/mod.rs

aria_kit/segmentation/
mod.rs
