beat_this/
mel.rs

1use anyhow::{anyhow, ensure, Result};
2
3use crate::runtime::{InferenceSession, 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 MelProcessor<S: InferenceSession> {
10    session: S,
11}
12
13impl<S: InferenceSession> MelProcessor<S> {
14    /// Wrap an already-loaded inference session for the mel spectrogram model.
15    pub fn new(session: S) -> Self {
16        Self { session }
17    }
18
19    /// Get a mutable reference to the underlying session.
20    pub fn session_mut(&mut self) -> &mut S {
21        &mut self.session
22    }
23
24    /// Compute mel spectrogram from mono PCM samples at 22050 Hz.
25    ///
26    /// Input: mono f32 samples (any length).
27    /// Output: Tensor with shape `[1, time_frames, 128]`.
28    ///
29    /// The number of time frames depends on sample count:
30    /// `time_frames ≈ samples.len() / 441` (hop_length=441 for 50 fps at 22050 Hz).
31    pub fn process(&mut self, samples: &[f32]) -> Result<Tensor> {
32        let input = Tensor {
33            shape: vec![1, samples.len()],
34            data: samples.to_vec(),
35        };
36
37        let mut outputs = self.session.run(&[("audio_pcm", &input)])?;
38
39        let mel = outputs
40            .remove("mel_spectrogram")
41            .ok_or_else(|| anyhow!("Model missing 'mel_spectrogram' output"))?;
42
43        ensure!(
44            mel.shape.len() == 3 && mel.shape[0] == 1 && mel.shape[2] == 128,
45            "Unexpected mel shape: {:?}",
46            mel.shape
47        );
48
49        Ok(mel)
50    }
51}
52
53/// Number of time frames in a mel spectrogram tensor.
54/// Assumes shape `[1, T, 128]`.
55pub fn num_frames(mel: &Tensor) -> usize {
56    mel.shape[1]
57}
beat_this/mel.rs

beat_this/
mel.rs
