audio_visualizer/spectrum/
png_file.rs

1//! Static spectrum analysis: print spectrum to PNG file.
2
3use crate::util::png::write_png_file_rgb_tuples;
4use std::collections::BTreeMap;
5use std::fs;
6use std::path::PathBuf;
7
8pub fn spectrum_static_png_visualize(
9    frequency_spectrum: &BTreeMap<u32, f32>,
10    directory: &str,
11    filename: &str,
12    highlighted_frequencies: &[f32],
13) {
14    // assert no NAN
15    assert!(
16        !frequency_spectrum.iter().any(|(_, f)| f.is_nan()),
17        "There are NAN-values in the spectrum!"
18    );
19
20    let image_width = 5000;
21    let image_height = 3000;
22
23    let mut rgb_img = vec![vec![(255, 255, 255); image_width]; image_height];
24
25    // find maximum for graphics scaling
26    let mut max = 0.0;
27    for mag in frequency_spectrum.values() {
28        if *mag > max {
29            max = *mag;
30        }
31    }
32
33    let x_step = image_width as f64 / frequency_spectrum.len() as f64;
34    for (i, (frequency, mag)) in frequency_spectrum.iter().enumerate() {
35        let mag = mag / max * image_height as f32;
36
37        let x = (i as f64 * x_step) as usize;
38
39        for j in 0..mag as usize {
40            let mut color = (0, 0, 0);
41
42            let highlight = highlighted_frequencies
43                .iter()
44                .any(|f| (*frequency as f32 - *f).abs() < 5.0);
45            if highlight {
46                color = (255, 0, 0);
47            }
48
49            // make it wider
50            if x > 2 && highlight {
51                rgb_img[image_height - 1 - j][x - 1] = color;
52                rgb_img[image_height - 1 - j][x - 2] = color;
53            }
54            rgb_img[image_height - 1 - j][x] = color;
55        }
56    }
57
58    if !fs::exists(directory).unwrap() {
59        fs::create_dir(directory).unwrap();
60    }
61    let mut path = PathBuf::new();
62    path.push(directory);
63    path.push(filename);
64    write_png_file_rgb_tuples(&path, &rgb_img);
65}
66
67#[cfg(test)]
68mod tests {
69    use super::*;
70    use crate::tests::testutil::TEST_OUT_DIR;
71
72    #[test]
73    fn test_visualize_sine_waves_spectrum() {
74        let mut spectrum = BTreeMap::new();
75        spectrum.insert(0, 0.0);
76        spectrum.insert(10, 5.0);
77        spectrum.insert(20, 20.0);
78        spectrum.insert(30, 40.0);
79        spectrum.insert(40, 80.0);
80        spectrum.insert(50, 120.0);
81        spectrum.insert(55, 130.0);
82        spectrum.insert(60, 140.0);
83        spectrum.insert(65, 130.0);
84        spectrum.insert(70, 120.0);
85        spectrum.insert(80, 80.0);
86        spectrum.insert(90, 40.0);
87        spectrum.insert(100, 20.0);
88        spectrum.insert(110, 5.0);
89        spectrum.insert(120, 0.0);
90        spectrum.insert(130, 0.0);
91
92        // Do FFT + get spectrum
93        spectrum_static_png_visualize(
94            &spectrum,
95            TEST_OUT_DIR,
96            "spectrum_60hz_peak_basic_visualization.png",
97            &[60.0],
98        );
99    }
100
101    #[allow(non_snake_case)]
102    #[test]
103    #[should_panic]
104    fn test_panic_on_NAN() {
105        let mut spectrum = BTreeMap::new();
106        spectrum.insert(0, f32::NAN);
107
108        spectrum_static_png_visualize(
109            &spectrum,
110            TEST_OUT_DIR,
111            "spectrum_60hz_peak_plotters_visualization_NAN.png",
112            &[],
113        );
114    }
115}