beat_this/
lib.rs

1pub mod audio;
2pub mod inference;
3pub mod mel;
4pub mod output;
5pub mod postprocessing;
6pub mod runtime;
7
8use std::path::Path;
9
10use anyhow::Result;
11
12pub use audio::{load_audio, AudioData};
13pub use inference::BeatInference;
14pub use mel::MelProcessor;
15pub use postprocessing::{BeatResult, PostProcessor};
16pub use runtime::{InferenceRuntime, InferenceSession, Tensor};
17
18/// Target sample rate expected by the mel spectrogram model.
19const TARGET_SAMPLE_RATE: u32 = 22050;
20
21/// High-level beat tracker composing the full pipeline.
22///
23/// Owns the mel spectrogram model, the beat inference model, and the
24/// post-processor. Generic over the inference session type, so it works
25/// with any backend (ort, rten, tract).
26pub struct BeatThis<S: InferenceSession> {
27    pub mel: MelProcessor<S>,
28    pub inference: BeatInference<S>,
29    pub post: PostProcessor,
30}
31
32impl<S: InferenceSession> BeatThis<S> {
33    /// Create a new beat tracker by loading both ONNX models via the given runtime.
34    ///
35    /// - `runtime`: any `InferenceRuntime` (e.g. `OrtRuntime::default()`)
36    /// - `mel_model_path`: path to the mel spectrogram ONNX model
37    /// - `beat_model_path`: path to the beat tracking ONNX model
38    pub fn new<R: InferenceRuntime<Session = S>>(
39        runtime: &R,
40        mel_model_path: &Path,
41        beat_model_path: &Path,
42    ) -> Result<Self> {
43        let mel_session = runtime.load_model(mel_model_path)?;
44        let beat_session = runtime.load_model(beat_model_path)?;
45
46        Ok(Self {
47            mel: MelProcessor::new(mel_session),
48            inference: BeatInference::new(beat_session),
49            post: PostProcessor::default(),
50        })
51    }
52
53    /// Run the full pipeline on an audio file.
54    ///
55    /// Loads the file, resamples to 22050 Hz mono, computes mel spectrogram,
56    /// runs beat inference, and post-processes into beat/downbeat timestamps.
57    pub fn process_file(&mut self, path: &Path) -> Result<BeatResult> {
58        let audio = load_audio(path, TARGET_SAMPLE_RATE)?;
59        self.process_audio(&audio.samples, audio.sample_rate)
60    }
61
62    /// Run the full pipeline on raw audio samples.
63    ///
64    /// The samples are resampled to 22050 Hz if `sample_rate` differs.
65    /// Input should be mono f32 PCM.
66    pub fn process_audio(&mut self, samples: &[f32], sample_rate: u32) -> Result<BeatResult> {
67        let samples = if sample_rate != TARGET_SAMPLE_RATE {
68            audio::resample(samples.to_vec(), sample_rate, TARGET_SAMPLE_RATE)?
69        } else {
70            samples.to_vec()
71        };
72
73        let mel = self.mel.process(&samples)?;
74        let (beat_logits, downbeat_logits) = self.inference.process(&mel)?;
75        self.post.process(&beat_logits, &downbeat_logits)
76    }
77}
beat_this/lib.rs

beat_this/
lib.rs
