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

beat_this/
lib.rs
