beat_this/
lib.rs

1mod audio;
2mod inference;
3mod mel;
4mod output;
5mod postprocessing;
6mod runtime;
7
8use std::path::Path;
9use std::time::Duration;
10
11use anyhow::Result;
12
13pub use audio::{load_audio, AudioData};
14pub use output::{beat_counts, calculate_bpm};
15#[cfg(feature = "ort")]
16pub use runtime::ort::OrtRuntime;
17pub use runtime::{rten::RtenRuntime, Model, Runtime, Tensor};
18
19use inference::BeatPredictor;
20use mel::MelExtractor;
21use postprocessing::PeakPicker;
22
23/// Target sample rate expected by the mel spectrogram model.
24const TARGET_SAMPLE_RATE: u32 = 22050;
25
26/// Full analysis result from the beat tracking pipeline.
27#[derive(Debug, Clone)]
28pub struct BeatAnalysis {
29    /// Beat times in seconds (sorted, deduplicated).
30    pub beats: Vec<f32>,
31    /// Downbeat times in seconds (sorted, deduplicated, snapped to nearest beat).
32    pub downbeats: Vec<f32>,
33    /// Mel spectrogram tensor with shape `[1, T, 128]` at 50 fps.
34    pub mel: Tensor,
35    /// Raw beat logits, one per spectrogram frame.
36    pub beat_logits: Vec<f32>,
37    /// Raw downbeat logits, one per spectrogram frame.
38    pub downbeat_logits: Vec<f32>,
39}
40
41/// High-level beat tracker composing the full pipeline.
42///
43/// Owns the mel spectrogram model, the beat prediction model, and the
44/// peak picker. Generic over the model type, so it works
45/// with any backend (ort, rten, tract).
46pub struct BeatThis<M: Model> {
47    mel: MelExtractor<M>,
48    predictor: BeatPredictor<M>,
49    peak_picker: PeakPicker,
50}
51
52/// Per-stage timing from [`BeatThis::analyze_audio_timed`].
53#[derive(Debug, Clone)]
54pub struct AnalysisTiming {
55    pub mel: Duration,
56    pub predict: Duration,
57    pub decode: Duration,
58}
59
60/// Analysis result with optional per-stage timing.
61#[derive(Debug, Clone)]
62pub struct TimedAnalysis {
63    pub analysis: BeatAnalysis,
64    pub timing: AnalysisTiming,
65}
66
67impl<M: Model> BeatThis<M> {
68    /// Create a new beat tracker by loading both ONNX models via the given runtime.
69    ///
70    /// - `runtime`: any `Runtime` (e.g. `OrtRuntime::default()`)
71    /// - `mel_model_path`: path to the mel spectrogram ONNX model
72    /// - `beat_model_path`: path to the beat tracking ONNX model
73    pub fn new<R: Runtime<Model = M>>(
74        runtime: &R,
75        mel_model_path: &Path,
76        beat_model_path: &Path,
77    ) -> Result<Self> {
78        let mel_model = runtime.load_model(mel_model_path)?;
79        let beat_model = runtime.load_model(beat_model_path)?;
80
81        Ok(Self {
82            mel: MelExtractor::new(mel_model),
83            predictor: BeatPredictor::new(beat_model),
84            peak_picker: PeakPicker::default(),
85        })
86    }
87
88    /// Create a beat tracker from pre-built models.
89    ///
90    /// Use this when you need separate runtimes for the mel and beat models
91    /// (e.g. one with profiling enabled).
92    pub fn from_models(mel_model: M, beat_model: M) -> Self {
93        Self {
94            mel: MelExtractor::new(mel_model),
95            predictor: BeatPredictor::new(beat_model),
96            peak_picker: PeakPicker::default(),
97        }
98    }
99
100    /// Get a mutable reference to the beat prediction model.
101    ///
102    /// Useful for runtime-specific operations like ending ORT profiling.
103    pub fn beat_model_mut(&mut self) -> &mut M {
104        self.predictor.model_mut()
105    }
106
107    /// Run the full pipeline on raw audio samples.
108    ///
109    /// The samples are resampled to 22050 Hz if `sample_rate` differs.
110    /// Input should be mono f32 PCM.
111    pub fn analyze_audio(&mut self, samples: &[f32], sample_rate: u32) -> Result<BeatAnalysis> {
112        Ok(self.analyze_audio_timed(samples, sample_rate)?.analysis)
113    }
114
115    /// Run the full pipeline on raw audio samples, returning per-stage timing.
116    pub fn analyze_audio_timed(
117        &mut self,
118        samples: &[f32],
119        sample_rate: u32,
120    ) -> Result<TimedAnalysis> {
121        let samples = if sample_rate != TARGET_SAMPLE_RATE {
122            audio::resample(samples.to_vec(), sample_rate, TARGET_SAMPLE_RATE)?
123        } else {
124            samples.to_vec()
125        };
126
127        let t = std::time::Instant::now();
128        let mel = self.mel.extract(&samples)?;
129        let mel_time = t.elapsed();
130
131        let t = std::time::Instant::now();
132        let (beat_logits, downbeat_logits) = self.predictor.predict(&mel)?;
133        let predict_time = t.elapsed();
134
135        let t = std::time::Instant::now();
136        let (beats, downbeats) = self.peak_picker.decode(&beat_logits, &downbeat_logits)?;
137        let decode_time = t.elapsed();
138
139        Ok(TimedAnalysis {
140            analysis: BeatAnalysis {
141                beats,
142                downbeats,
143                mel,
144                beat_logits,
145                downbeat_logits,
146            },
147            timing: AnalysisTiming {
148                mel: mel_time,
149                predict: predict_time,
150                decode: decode_time,
151            },
152        })
153    }
154
155    /// Run the full pipeline on an audio file.
156    ///
157    /// Loads the file, resamples to 22050 Hz mono, computes mel spectrogram,
158    /// runs beat prediction, and decodes into beat/downbeat timestamps.
159    pub fn analyze_file(&mut self, path: &Path) -> Result<BeatAnalysis> {
160        let audio = load_audio(path, TARGET_SAMPLE_RATE)?;
161        self.analyze_audio(&audio.samples, audio.sample_rate)
162    }
163}
beat_this/lib.rs

beat_this/
lib.rs
