Files
rustytorch/crates/models/rtx-csm/src/ser.rs
T
2026-05-07 16:30:04 +00:00

395 lines
14 KiB
Rust

//! Speech Emotion Recognition — prosody-rule baseline.
//!
//! No ML model. Extracts F0 (autocorrelation) + RMS energy + voicing ratio
//! and maps to a 5-bucket emotion label via hand-tuned thresholds. Crude
//! but useful as a default for the data-prep `--auto-emotion-tag` path
//! (Phase 13.2) when no labeled data exists.
//!
//! The thresholds are gut-feel; real-world tuning will want either a corpus
//! pass for z-score normalization or a swap to a real classifier
//! (emotion2vec_plus_base port). The [`EmotionDetector`] trait makes that
//! swap a one-line change.
//!
//! Honest scope: this is a placeholder so the pipeline is complete.
//! Don't ship a product that depends on these labels being right.
use crate::error::Result;
/// Emotion buckets used as tags by the data-prep + reactive-emotion paths.
///
/// The five `Neutral / Calm / Sad / Angry / Excited` buckets are the
/// original Phase 13.3 prosody-rule output set. Phase 13.8's emotion2vec
/// port emits the four extra raw classes (`Disgusted`, `Fearful`, `Happy`,
/// `Surprised`) directly — folding them into the original 5 collapses
/// emotion2vec's resolution to a single `[excited]` tag in practice
/// (verified empirically: motivational speech, technical talks, and movie
/// dialogue ALL bucketed to `[excited]` under the old 9→5 fold).
///
/// Tags are bracket-wrapped lowercase and slot directly into Phase 12.2's
/// `GenerateOptions::emotion_hint`. `Unk` is emotion2vec's `<unk>` class.
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize)]
pub enum EmotionLabel {
Neutral,
Calm,
Sad,
Angry,
Excited,
Disgusted,
Fearful,
Happy,
Surprised,
Unk,
}
impl EmotionLabel {
pub fn as_tag(&self) -> &'static str {
match self {
EmotionLabel::Neutral => "[neutral]",
EmotionLabel::Calm => "[calm]",
EmotionLabel::Sad => "[sad]",
EmotionLabel::Angry => "[angry]",
EmotionLabel::Excited => "[excited]",
EmotionLabel::Disgusted => "[disgusted]",
EmotionLabel::Fearful => "[fearful]",
EmotionLabel::Happy => "[happy]",
EmotionLabel::Surprised => "[surprised]",
EmotionLabel::Unk => "[unk]",
}
}
}
/// Trait for any classifier that maps a raw 16 kHz mono sample slice to
/// an [`EmotionLabel`]. Implemented today by [`ProsodyDetector`]; in the
/// future a candle-native `Emotion2VecDetector` would slot in here.
pub trait EmotionDetector {
fn classify(&self, samples_16k: &[f32]) -> Result<EmotionLabel>;
}
/// Per-utterance prosodic features used by [`ProsodyDetector`]. Surfaced
/// so callers can introspect what drove a given label.
#[derive(Debug, Clone)]
pub struct ProsodyFeatures {
pub rms_mean: f32,
pub f0_mean_hz: f32,
pub f0_std_hz: f32,
pub voiced_ratio: f32,
}
/// Cutoff thresholds for the 5-bucket classifier. All thresholds are
/// absolute (not corpus-normalized). For a real deployment you'd
/// z-score-normalize within a corpus pass and compare to means; for the
/// out-of-the-box default these mid-range values are okay-ish for typical
/// 16-bit speech material at speaker-level RMS.
#[derive(Debug, Clone)]
pub struct ProsodyThresholds {
/// Above this RMS = "high energy".
pub high_energy: f32,
/// Below this RMS = "low energy".
pub low_energy: f32,
/// Above this F0 std = "high pitch variance" (expressive).
pub high_f0_std: f32,
/// Below this F0 std = "low pitch variance" (flat).
pub low_f0_std: f32,
/// Below this voiced fraction → fall back to Neutral (mostly unvoiced
/// or noisy input — pitch stats are unreliable).
pub min_voiced_ratio: f32,
}
impl Default for ProsodyThresholds {
fn default() -> Self {
Self {
high_energy: 0.15,
low_energy: 0.05,
high_f0_std: 35.0,
low_f0_std: 15.0,
min_voiced_ratio: 0.30,
}
}
}
/// Pure-DSP prosody-based emotion classifier.
///
/// Pipeline per utterance:
/// 1. Frame the audio at `frame_samples` with `hop_samples` hop
/// 2. Per frame: RMS, plus F0 via autocorrelation (clamped to
/// `min_f0_hz..=max_f0_hz`); flag voiced if peak/total ratio >
/// `voiced_threshold`
/// 3. Aggregate: rms_mean across all frames, f0_mean / f0_std across
/// voiced frames, voiced_ratio = voiced_frames / total_frames
/// 4. Apply [`ProsodyThresholds`] to bucket
#[derive(Debug, Clone)]
pub struct ProsodyDetector {
pub frame_samples: usize,
pub hop_samples: usize,
pub min_f0_hz: f32,
pub max_f0_hz: f32,
pub voiced_threshold: f32,
pub sample_rate: u32,
pub thresholds: ProsodyThresholds,
}
impl Default for ProsodyDetector {
fn default() -> Self {
// 25 ms frame, 10 ms hop @ 16 kHz — standard speech analysis.
Self {
frame_samples: 400,
hop_samples: 160,
// Voice F0 range. 65 Hz catches deep male; 400 Hz covers
// expressive female / shouting.
min_f0_hz: 65.0,
max_f0_hz: 400.0,
voiced_threshold: 0.4,
sample_rate: 16_000,
thresholds: ProsodyThresholds::default(),
}
}
}
impl ProsodyDetector {
/// Compute the prosodic features without classifying. Useful for
/// debugging or for callers that want to do their own bucketing.
pub fn features(&self, samples: &[f32]) -> ProsodyFeatures {
if samples.len() < self.frame_samples {
return ProsodyFeatures {
rms_mean: 0.0,
f0_mean_hz: 0.0,
f0_std_hz: 0.0,
voiced_ratio: 0.0,
};
}
let mut rms_acc = 0.0f64;
let mut rms_n = 0u32;
let mut voiced_f0s: Vec<f32> = Vec::new();
let mut total_frames = 0u32;
let mut t = 0;
let min_lag = (self.sample_rate as f32 / self.max_f0_hz).floor() as usize;
let max_lag = (self.sample_rate as f32 / self.min_f0_hz).ceil() as usize;
while t + self.frame_samples <= samples.len() {
let frame = &samples[t..t + self.frame_samples];
// RMS
let sumsq: f64 = frame.iter().map(|x| (*x as f64) * (*x as f64)).sum();
let rms = (sumsq / frame.len() as f64).sqrt() as f32;
rms_acc += rms as f64;
rms_n += 1;
// F0 via autocorrelation
if let Some((f0, voiced)) = autocorr_f0(
frame,
self.sample_rate,
min_lag,
max_lag,
self.voiced_threshold,
) && voiced
&& f0 >= self.min_f0_hz
&& f0 <= self.max_f0_hz
{
voiced_f0s.push(f0);
}
total_frames += 1;
t += self.hop_samples;
}
let rms_mean = if rms_n > 0 {
(rms_acc / rms_n as f64) as f32
} else {
0.0
};
let voiced_ratio = if total_frames > 0 {
voiced_f0s.len() as f32 / total_frames as f32
} else {
0.0
};
let (f0_mean, f0_std) = if voiced_f0s.is_empty() {
(0.0, 0.0)
} else {
let mean = voiced_f0s.iter().sum::<f32>() / voiced_f0s.len() as f32;
let var = voiced_f0s.iter().map(|f| (*f - mean).powi(2)).sum::<f32>()
/ voiced_f0s.len() as f32;
(mean, var.sqrt())
};
ProsodyFeatures {
rms_mean,
f0_mean_hz: f0_mean,
f0_std_hz: f0_std,
voiced_ratio,
}
}
pub fn classify_features(&self, f: &ProsodyFeatures) -> EmotionLabel {
let t = &self.thresholds;
if f.voiced_ratio < t.min_voiced_ratio {
return EmotionLabel::Neutral;
}
let high_e = f.rms_mean > t.high_energy;
let low_e = f.rms_mean < t.low_energy;
let high_var = f.f0_std_hz > t.high_f0_std;
let low_var = f.f0_std_hz < t.low_f0_std;
match (high_e, low_e, high_var, low_var) {
(true, _, true, _) => EmotionLabel::Excited,
(true, _, _, true) => EmotionLabel::Angry,
(_, true, _, true) => EmotionLabel::Sad,
(_, true, true, _) => EmotionLabel::Calm,
_ => EmotionLabel::Neutral,
}
}
}
impl EmotionDetector for ProsodyDetector {
fn classify(&self, samples_16k: &[f32]) -> Result<EmotionLabel> {
let f = self.features(samples_16k);
Ok(self.classify_features(&f))
}
}
/// Single-frame F0 estimation via normalized autocorrelation.
///
/// Returns `(f0_hz, voiced)` where `voiced` is true if the autocorr peak
/// at the best lag was above `voiced_threshold * autocorr[0]`. Falls back
/// to None when the frame is too quiet to estimate.
fn autocorr_f0(
frame: &[f32],
sample_rate: u32,
min_lag: usize,
max_lag: usize,
voiced_threshold: f32,
) -> Option<(f32, bool)> {
let n = frame.len();
if max_lag >= n {
return None;
}
// Mean-remove + normalize the frame to make the autocorr peak ratio
// independent of input gain.
let mean = frame.iter().sum::<f32>() / n as f32;
let centered: Vec<f32> = frame.iter().map(|x| x - mean).collect();
let r0: f32 = centered.iter().map(|x| x * x).sum();
if r0 < 1e-6 {
return None; // silence
}
let mut best_lag = 0usize;
let mut best_corr = 0.0f32;
for lag in min_lag..=max_lag {
let mut corr = 0.0f32;
for i in 0..(n - lag) {
corr += centered[i] * centered[i + lag];
}
if corr > best_corr {
best_corr = corr;
best_lag = lag;
}
}
if best_lag == 0 {
return None;
}
let f0 = sample_rate as f32 / best_lag as f32;
let voiced = (best_corr / r0) > voiced_threshold;
Some((f0, voiced))
}
#[cfg(test)]
mod tests {
use super::*;
fn sine(freq: f32, dur_s: f32, sr: u32, amp: f32) -> Vec<f32> {
let n = (dur_s * sr as f32) as usize;
(0..n)
.map(|i| amp * (2.0 * std::f32::consts::PI * freq * i as f32 / sr as f32).sin())
.collect()
}
#[test]
fn autocorr_recovers_pure_tone_pitch() {
let sr = 16_000u32;
let frame = sine(220.0, 0.025, sr, 0.5); // 25 ms @ 220 Hz
let min_lag = (sr as f32 / 400.0).floor() as usize;
let max_lag = (sr as f32 / 65.0).ceil() as usize;
let (f0, voiced) = autocorr_f0(&frame, sr, min_lag, max_lag, 0.4).unwrap();
assert!(voiced, "pure tone should be voiced");
assert!((f0 - 220.0).abs() < 5.0, "expected ~220 Hz, got {f0}");
}
#[test]
fn autocorr_returns_none_on_silence() {
let frame = vec![0.0f32; 400];
let r = autocorr_f0(&frame, 16_000, 40, 250, 0.4);
assert!(r.is_none());
}
#[test]
fn prosody_low_energy_flat_pitch_is_sad() {
// Quiet, fixed-pitch tone — should bucket as Sad (low E + low variance).
let det = ProsodyDetector::default();
let audio = sine(120.0, 1.0, 16_000, 0.02);
let f = det.features(&audio);
assert!(f.voiced_ratio > 0.3, "should be voiced enough");
assert!(
f.rms_mean < det.thresholds.low_energy,
"should be low energy: {}",
f.rms_mean
);
assert!(
f.f0_std_hz < det.thresholds.low_f0_std,
"pure tone should be flat: {}",
f.f0_std_hz
);
assert_eq!(det.classify_features(&f), EmotionLabel::Sad);
}
#[test]
fn prosody_silence_falls_back_to_neutral() {
let det = ProsodyDetector::default();
let audio = vec![0.0f32; 16_000];
let label = det.classify(&audio).unwrap();
assert_eq!(label, EmotionLabel::Neutral);
}
#[test]
fn prosody_high_energy_high_variance_is_excited() {
// Synthesize a loud chirp (linear pitch sweep) — high energy +
// high pitch variance should bucket as Excited.
let sr = 16_000u32;
let dur = 1.0;
let n = (dur * sr as f32) as usize;
let amp = 0.4;
let f_start = 120.0;
let f_end = 280.0;
let mut audio = Vec::with_capacity(n);
let mut phase = 0.0f32;
for i in 0..n {
let t = i as f32 / sr as f32;
let f = f_start + (f_end - f_start) * (t / dur);
phase += 2.0 * std::f32::consts::PI * f / sr as f32;
audio.push(amp * phase.sin());
}
let det = ProsodyDetector::default();
let f = det.features(&audio);
// We expect this to be voiced with high variance and high enough energy.
assert!(f.voiced_ratio > 0.5, "voiced_ratio={}", f.voiced_ratio);
assert!(
f.rms_mean > det.thresholds.high_energy,
"rms={}",
f.rms_mean
);
assert!(
f.f0_std_hz > det.thresholds.high_f0_std,
"f0_std={}",
f.f0_std_hz
);
assert_eq!(det.classify_features(&f), EmotionLabel::Excited);
}
#[test]
fn emotion_label_tags_match_phase_12_2_format() {
// The tags need to round-trip through GenerateOptions::emotion_hint
// verbatim, so guard the format: bracket-enclosed lowercase.
for lab in [
EmotionLabel::Neutral,
EmotionLabel::Calm,
EmotionLabel::Sad,
EmotionLabel::Angry,
EmotionLabel::Excited,
] {
let t = lab.as_tag();
assert!(t.starts_with('[') && t.ends_with(']'), "tag format: {t}");
assert_eq!(t.to_lowercase(), t, "tag should be lowercase: {t}");
}
}
}