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