audio_processing.rs (9367B)
1 use std::sync::LazyLock; 2 3 use crate::constants::{ 4 BPM_HOP_SIZE, BPM_MAX, BPM_MIN, BPM_OFFSET, BPM_WINDOW_SIZE, CHROMA_BINS, KEY_HOP_SIZE, 5 KEY_WINDOW_SIZE, SAMPLE_RATE, TUNING_PRECISION, 6 }; 7 use anyhow::Result; 8 use bliss_audio::chroma::bench::{chroma_stft, estimate_tuning}; 9 use bliss_audio::utils::bench::stft; 10 use bliss_audio_aubio_rs::{OnsetMode, Tempo}; 11 use ndarray::Array1; 12 13 /// Calculate the BPM of a track. 14 /// 15 /// # Panics 16 /// 17 /// Panics if the internal tempo object cannot be created or if a single tempo 18 /// calculation fails. 19 #[must_use] 20 pub fn calculate_bpm(samples: &[f32]) -> f32 { 21 let window_size = BPM_WINDOW_SIZE; 22 let hop_size = BPM_HOP_SIZE; 23 let mut tempo = Tempo::new(OnsetMode::SpecFlux, window_size, hop_size, SAMPLE_RATE) 24 .expect("Failed to create tempo object"); 25 let mut bpms = Vec::with_capacity(samples.len() / hop_size); 26 27 for chunk in samples.chunks(hop_size) { 28 let mut padded = chunk.to_vec(); 29 if padded.len() < hop_size { 30 padded.resize(hop_size, 0.0); 31 } 32 if tempo.do_result(&padded).expect("Failed to calculate tempo") > 0.0 { 33 bpms.push(tempo.get_bpm()); 34 } 35 } 36 37 if bpms.is_empty() { 38 return 0.0; 39 } 40 41 bpms.sort_by(f32::total_cmp); 42 let median_bpm = bpms[bpms.len() / 2]; 43 44 normalize_bpm(median_bpm) 45 } 46 47 fn normalize_bpm(mut bpm: f32) -> f32 { 48 if bpm > 0.0 { 49 bpm += BPM_OFFSET; 50 } 51 52 if bpm > 0.0 { 53 while bpm > BPM_MAX { 54 bpm /= 2.0; 55 } 56 while bpm < BPM_MIN { 57 bpm *= 2.0; 58 } 59 } 60 61 bpm 62 } 63 64 /// Calculate the key of a track. 65 /// 66 /// # Errors 67 /// 68 /// Returns an error if the key cannot be calculated. 69 pub fn calculate_key(samples: &[f32]) -> Result<String> { 70 let global_chroma = calculate_global_chroma(samples)?; 71 let (best_key_idx, best_mode) = find_best_key(&global_chroma); 72 let key_names = [ 73 "C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B", 74 ]; 75 Ok(format!("{} {}", key_names[best_key_idx], best_mode)) 76 } 77 78 #[allow(clippy::cast_possible_truncation)] 79 fn calculate_global_chroma(samples: &[f32]) -> Result<Array1<f64>> { 80 let window_size = KEY_WINDOW_SIZE; 81 let hop_size = KEY_HOP_SIZE; 82 83 let mut spectrum = stft(samples, window_size, hop_size); 84 85 let tuning = estimate_tuning( 86 SAMPLE_RATE, 87 &spectrum, 88 window_size, 89 TUNING_PRECISION, 90 CHROMA_BINS as _, 91 )?; 92 93 let n_chroma = CHROMA_BINS as u32; 94 let chroma = chroma_stft(SAMPLE_RATE, &mut spectrum, window_size, n_chroma, tuning)?; 95 96 let (n_bins, n_frames) = (chroma.shape()[0], chroma.shape()[1]); 97 let mut result = vec![0.0f64; n_bins]; 98 for i in 0..n_bins { 99 for j in 0..n_frames { 100 result[i] += chroma[[i, j]]; 101 } 102 } 103 104 Ok(Array1::from_vec(result)) 105 } 106 107 static MAJOR_ROTATIONS: LazyLock<Vec<Array1<f64>>> = LazyLock::new(|| { 108 let base = ndarray::arr1(&[ 109 6.35, 2.23, 3.48, 2.33, 4.38, 4.09, 2.52, 5.19, 2.39, 3.66, 2.29, 2.88, 110 ]); 111 (0..CHROMA_BINS).map(|i| rotate_array(&base, i)).collect() 112 }); 113 114 static MINOR_ROTATIONS: LazyLock<Vec<Array1<f64>>> = LazyLock::new(|| { 115 let base = ndarray::arr1(&[ 116 6.33, 2.68, 3.52, 5.38, 2.60, 3.53, 2.54, 4.75, 3.98, 2.69, 3.34, 3.17, 117 ]); 118 (0..CHROMA_BINS).map(|i| rotate_array(&base, i)).collect() 119 }); 120 121 fn find_best_key(global_chroma: &Array1<f64>) -> (usize, &'static str) { 122 let mut max_corr = -1.0; 123 let mut best_key_idx = 0; 124 let mut best_mode = "Major"; 125 let (major_mean, major_stddev) = *MAJOR_PROFILE_STATS; 126 let (minor_mean, minor_stddev) = *MINOR_PROFILE_STATS; 127 128 for i in 0..CHROMA_BINS { 129 let corr_major = 130 pearson_correlation(global_chroma, &MAJOR_ROTATIONS[i], major_mean, major_stddev); 131 if corr_major > max_corr { 132 max_corr = corr_major; 133 best_key_idx = i; 134 best_mode = "Major"; 135 } 136 137 let corr_minor = 138 pearson_correlation(global_chroma, &MINOR_ROTATIONS[i], minor_mean, minor_stddev); 139 if corr_minor > max_corr { 140 max_corr = corr_minor; 141 best_key_idx = i; 142 best_mode = "Minor"; 143 } 144 } 145 146 (best_key_idx, best_mode) 147 } 148 149 fn rotate_array(arr: &Array1<f64>, shift: usize) -> Array1<f64> { 150 let n = arr.len(); 151 Array1::from_iter((0..n).map(|j| arr[(j + n - shift) % n])) 152 } 153 154 /// Pre-computed profile stats: (mean, stddev) for major and minor profiles. 155 /// Since rotation is a cyclic permutation, all rotations share the same mean and stddev. 156 static MAJOR_PROFILE_STATS: LazyLock<(f64, f64)> = LazyLock::new(|| { 157 let v = &MAJOR_ROTATIONS[0]; 158 let mean = v.mean().expect("Failed to calculate mean"); 159 let denom: f64 = v.iter().map(|x| (x - mean).powi(2)).sum(); 160 (mean, denom.sqrt()) 161 }); 162 163 static MINOR_PROFILE_STATS: LazyLock<(f64, f64)> = LazyLock::new(|| { 164 let v = &MINOR_ROTATIONS[0]; 165 let mean = v.mean().expect("Failed to calculate mean"); 166 let denom: f64 = v.iter().map(|x| (x - mean).powi(2)).sum(); 167 (mean, denom.sqrt()) 168 }); 169 170 fn pearson_correlation( 171 global_chroma: &Array1<f64>, 172 profile: &Array1<f64>, 173 profile_mean: f64, 174 profile_stddev: f64, 175 ) -> f64 { 176 let mean1 = global_chroma.mean().expect("Failed to calculate mean"); 177 let num: f64 = global_chroma 178 .iter() 179 .zip(profile.iter()) 180 .map(|(x, y)| (x - mean1) * (y - profile_mean)) 181 .sum(); 182 let den1: f64 = global_chroma.iter().map(|x| (x - mean1).powi(2)).sum(); 183 if den1 == 0.0 || profile_stddev == 0.0 { 184 0.0 185 } else { 186 num / (den1.sqrt() * profile_stddev) 187 } 188 } 189 190 #[cfg(test)] 191 mod tests { 192 use super::*; 193 194 #[test] 195 fn normalize_bpm_applies_offset_within_range() { 196 // 120 + (-1.15) stays within [80, 160]. 197 assert!((normalize_bpm(120.0) - 118.85).abs() < 1e-4); 198 } 199 200 #[test] 201 fn normalize_bpm_halves_when_too_fast() { 202 // 320 - 1.15 = 318.85 -> /2 = 159.425 (within range). 203 assert!((normalize_bpm(320.0) - 159.425).abs() < 1e-3); 204 } 205 206 #[test] 207 fn normalize_bpm_doubles_when_too_slow() { 208 // 40 - 1.15 = 38.85 -> *2 = 77.7 -> *2 = 155.4 (within range). 209 assert!((normalize_bpm(40.0) - 155.4).abs() < 1e-3); 210 } 211 212 #[test] 213 fn normalize_bpm_keeps_zero() { 214 assert_eq!(normalize_bpm(0.0), 0.0); 215 } 216 217 #[test] 218 fn normalized_bpm_lands_in_range_for_many_inputs() { 219 // A detected tempo below ~2 BPM isn't a meaningful value; the offset 220 // can push such inputs non-positive, bypassing octave normalization. 221 for raw in 2..=2000 { 222 let bpm = normalize_bpm(raw as f32); 223 assert!( 224 (BPM_MIN..=BPM_MAX).contains(&bpm), 225 "normalize_bpm({raw}) = {bpm} outside [{BPM_MIN}, {BPM_MAX}]" 226 ); 227 } 228 } 229 230 #[test] 231 fn rotate_array_shifts_right() { 232 let arr = ndarray::arr1(&[0.0, 1.0, 2.0, 3.0]); 233 assert_eq!(rotate_array(&arr, 0), arr); 234 assert_eq!(rotate_array(&arr, 1), ndarray::arr1(&[3.0, 0.0, 1.0, 2.0])); 235 assert_eq!(rotate_array(&arr, 2), ndarray::arr1(&[2.0, 3.0, 0.0, 1.0])); 236 // A full rotation returns the original. 237 assert_eq!(rotate_array(&arr, 4), arr); 238 } 239 240 #[test] 241 fn pearson_correlation_of_identical_vectors_is_one() { 242 let v = ndarray::arr1(&[1.0_f64, 2.0, 3.0, 4.0, 5.0]); 243 let mean = v.mean().unwrap(); 244 let stddev: f64 = v.iter().map(|x| (x - mean).powi(2)).sum::<f64>().sqrt(); 245 let corr = pearson_correlation(&v, &v, mean, stddev); 246 assert!((corr - 1.0).abs() < 1e-9, "expected 1.0, got {corr}"); 247 } 248 249 #[test] 250 fn pearson_correlation_of_constant_input_is_zero() { 251 // A flat chroma has zero variance, so correlation is defined as 0. 252 let v = ndarray::arr1(&[2.0; CHROMA_BINS]); 253 let (mean, stddev) = *MAJOR_PROFILE_STATS; 254 assert_eq!(pearson_correlation(&v, &MAJOR_ROTATIONS[0], mean, stddev), 0.0); 255 } 256 257 #[test] 258 fn profile_stats_share_mean_and_stddev_across_rotations() { 259 // Rotation is a permutation, so every rotation has identical stats. 260 let (mean, stddev) = *MAJOR_PROFILE_STATS; 261 for rot in MAJOR_ROTATIONS.iter() { 262 let m = rot.mean().unwrap(); 263 let s: f64 = rot.iter().map(|x| (x - m).powi(2)).sum::<f64>().sqrt(); 264 assert!((m - mean).abs() < 1e-9); 265 assert!((s - stddev).abs() < 1e-9); 266 } 267 assert!(stddev > 0.0); 268 } 269 270 #[test] 271 fn find_best_key_recovers_each_major_profile() { 272 let key_names = [ 273 "C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B", 274 ]; 275 for i in 0..CHROMA_BINS { 276 // A chroma that is exactly the i-th major profile must resolve to 277 // that key in Major mode (self-correlation is the global maximum). 278 let (idx, mode) = find_best_key(&MAJOR_ROTATIONS[i]); 279 assert_eq!(idx, i, "{} major misidentified", key_names[i]); 280 assert_eq!(mode, "Major"); 281 } 282 } 283 284 #[test] 285 fn find_best_key_recovers_each_minor_profile() { 286 for i in 0..CHROMA_BINS { 287 let (idx, mode) = find_best_key(&MINOR_ROTATIONS[i]); 288 assert_eq!(idx, i); 289 assert_eq!(mode, "Minor"); 290 } 291 } 292 }