beat_this/
mel.rs

1use anyhow::{anyhow, ensure, Result};
2
3use crate::runtime::{Model, Tensor};
4
5/// Computes log-mel spectrograms via an ONNX model.
6///
7/// The model takes raw PCM audio and returns a mel spectrogram,
8/// guaranteeing exact numerical parity with the Python training pipeline.
9pub struct MelExtractor<M: Model> {
10    model: M,
11}
12
13impl<M: Model> MelExtractor<M> {
14    /// Wrap an already-loaded model for mel spectrogram extraction.
15    pub fn new(model: M) -> Self {
16        Self { model }
17    }
18
19    /// Extract mel spectrogram from mono PCM samples at 22050 Hz.
20    ///
21    /// Input: mono f32 samples (any length).
22    /// Output: Tensor with shape `[1, time_frames, 128]`.
23    ///
24    /// The number of time frames depends on sample count:
25    /// `time_frames ≈ samples.len() / 441` (hop_length=441 for 50 fps at 22050 Hz).
26    pub fn extract(&mut self, samples: &[f32]) -> Result<Tensor> {
27        let input = Tensor {
28            shape: vec![1, samples.len()],
29            data: samples.to_vec(),
30        };
31
32        let mut outputs = self.model.run(&[("audio_pcm", &input)])?;
33
34        let mel = outputs
35            .remove("mel_spectrogram")
36            .ok_or_else(|| anyhow!("Model missing 'mel_spectrogram' output"))?;
37
38        ensure!(
39            mel.shape.len() == 3 && mel.shape[0] == 1 && mel.shape[2] == 128,
40            "Unexpected mel shape: {:?}",
41            mel.shape
42        );
43
44        Ok(mel)
45    }
46}
beat_this/mel.rs

beat_this/
mel.rs
