Skip to main content

sound_to_patch/report/
spectrogram.rs

1//! The spectrogram image: two log-frequency panels of a 10 s passage.
2//!
3//! The image gives the reader coarse shape only; the numbers live in
4//! `layers.json`. The full panel uses 4096-point frames from 20 Hz to
5//! 20 kHz. The zoom panel uses 16384-point frames from 20 to 500 Hz, where
6//! 4096-point bins of about 11 Hz put a 41 Hz fundamental and its second
7//! harmonic only 4 bins apart.
8
9use font8x8::UnicodeFonts as _;
10
11use crate::audio::SampleRate;
12use crate::cast;
13use crate::dsp::{ActiveFrames, Stft, StftParams, Window};
14use crate::units::{Decibels, Hz, Seconds};
15
16/// The image width in pixels. The long edge stays below 1568 pixels, where
17/// the image reader of the hosted model downscales.
18pub const WIDTH: usize = 1500;
19/// The image height in pixels.
20pub const HEIGHT: usize = 512;
21/// The length of the passage the image shows.
22pub const PASSAGE_S: Seconds = Seconds::from_const(10.0);
23/// The colour range below the stem peak.
24pub const RANGE_DB: Decibels = Decibels::from_const(80.0);
25
26/// One panel: its frame length and its frequency range.
27#[derive(Clone, Copy, Debug, PartialEq)]
28pub struct PanelSpec {
29    /// The STFT frame length.
30    pub frame_len: usize,
31    /// The lowest frequency, at the bottom edge.
32    pub lo: Hz,
33    /// The highest frequency, at the top edge.
34    pub hi: Hz,
35}
36
37/// The two panels, top first.
38pub const PANELS: [PanelSpec; 2] = [
39    PanelSpec {
40        frame_len: 4096,
41        lo: Hz::from_const(20.0),
42        hi: Hz::from_const(20_000.0),
43    },
44    PanelSpec {
45        frame_len: 16384,
46        lo: Hz::from_const(20.0),
47        hi: Hz::from_const(500.0),
48    },
49];
50
51const LEFT: usize = 48;
52const RIGHT: usize = 60;
53const TITLE: usize = 12;
54const PANEL_HEIGHT: usize = 232;
55/// The plot width in pixels: one STFT frame per column.
56pub const PLOT_WIDTH: usize = WIDTH - LEFT - RIGHT;
57/// The plot height of each panel in pixels: one frequency band per row.
58pub const PLOT_HEIGHT: usize = PANEL_HEIGHT;
59const AXIS_TOP: usize = 2 * (TITLE + PANEL_HEIGHT);
60const BACKGROUND: [u8; 3] = [0, 0, 0];
61const INK: [u8; 3] = [255, 255, 255];
62const FRAME: [u8; 3] = [128, 128, 128];
63
64/// The level matrix of one panel in dB below the stem peak, before the
65/// colour map.
66#[derive(Clone, Debug, PartialEq)]
67pub struct PanelDb {
68    spec: PanelSpec,
69    db: Vec<f64>,
70}
71
72impl PanelDb {
73    /// Returns the panel's frame length and range.
74    #[must_use]
75    pub fn spec(&self) -> PanelSpec {
76        self.spec
77    }
78
79    /// Returns the levels row by row, top row first, [`PLOT_WIDTH`] per
80    /// row and [`PLOT_HEIGHT`] rows, each from `-RANGE_DB` to 0.
81    #[must_use]
82    pub fn db(&self) -> &[f64] {
83        &self.db
84    }
85}
86
87/// The spectrogram of a passage.
88#[derive(Clone, Debug, PartialEq)]
89pub struct Spectrogram {
90    start: Seconds,
91    span: Seconds,
92    panels: Vec<PanelDb>,
93}
94
95impl Spectrogram {
96    /// Computes both panels of a mono signal.
97    ///
98    /// The passage is the [`PASSAGE_S`] window that holds the most active
99    /// 4096-point frames, the earliest on a tie, or the whole signal when
100    /// it is shorter. 0 dB is the loudest bin above DC over the whole
101    /// signal at the panel's frame length, hop `N / 4`, so both panels
102    /// share one scale up to the frame length's resolution.
103    #[must_use]
104    pub fn of(samples: &[f64], rate: SampleRate) -> Self {
105        let (first, len) = passage(samples, rate);
106        let panels = PANELS
107            .iter()
108            .map(|&spec| PanelDb {
109                spec,
110                db: panel_db(samples, rate, spec, first, len),
111            })
112            .collect();
113        Self {
114            start: rate.seconds(first),
115            span: rate.seconds(len),
116            panels,
117        }
118    }
119
120    /// Returns the start of the passage.
121    #[must_use]
122    pub fn start(&self) -> Seconds {
123        self.start
124    }
125
126    /// Returns the length of the passage.
127    #[must_use]
128    pub fn span(&self) -> Seconds {
129        self.span
130    }
131
132    /// Returns the panels, top first.
133    #[must_use]
134    pub fn panels(&self) -> &[PanelDb] {
135        &self.panels
136    }
137
138    /// Returns the RGB pixels, row by row, [`WIDTH`] by [`HEIGHT`].
139    #[must_use]
140    pub fn pixels(&self) -> Vec<u8> {
141        let mut canvas = Canvas::new();
142        for (index, panel) in self.panels.iter().enumerate() {
143            let top = index * (TITLE + PANEL_HEIGHT);
144            let spec = panel.spec;
145            canvas.text(
146                LEFT,
147                top + 2,
148                &format!(
149                    "N = {}, {} Hz to {} Hz, dB below the stem peak",
150                    spec.frame_len,
151                    hz_label(spec.lo.get()),
152                    hz_label(spec.hi.get())
153                ),
154            );
155            let plot_top = top + TITLE;
156            for (row, levels) in panel.db.chunks(PLOT_WIDTH).enumerate() {
157                for (column, &db) in levels.iter().enumerate() {
158                    canvas.set(LEFT + column, plot_top + row, colour(db));
159                }
160            }
161            canvas.outline(LEFT - 1, plot_top - 1, PLOT_WIDTH + 2, PANEL_HEIGHT + 2);
162            for &tick in frequency_ticks(spec) {
163                let row = frequency_row(spec, tick);
164                let y = plot_top + row;
165                canvas.hline(LEFT - 4, y, 3);
166                let label = hz_label(tick);
167                let x = LEFT - 5 - 8 * label.len();
168                canvas.text(x, y.saturating_sub(3), &label);
169            }
170        }
171        canvas.text(0, 2, "Hz");
172        self.time_axis(&mut canvas);
173        colour_bar(&mut canvas);
174        canvas.pixels
175    }
176
177    fn time_axis(&self, canvas: &mut Canvas) {
178        let span = self.span.get();
179        let start = self.start.get();
180        let step = if span >= 2.0 { 1.0 } else { 0.1 };
181        let first = (start / step).ceil();
182        let mut index = 0.0;
183        while (first + index) * step <= start + span + 1e-9 {
184            let time = (first + index) * step;
185            index += 1.0;
186            let offset = if span > 0.0 {
187                (time - start) / span
188            } else {
189                0.0
190            };
191            let column = cast::round_to_usize(offset * cast::usize_to_f64(PLOT_WIDTH - 1))
192                .unwrap_or(0)
193                .min(PLOT_WIDTH - 1);
194            let x = LEFT + column;
195            canvas.vline(x, AXIS_TOP, 4);
196            let label = if step < 1.0 {
197                format!("{time:.1}")
198            } else {
199                format!("{time:.0}")
200            };
201            let width = 8 * label.len();
202            canvas.text(
203                x.saturating_sub(width / 2).min(WIDTH - RIGHT - width),
204                AXIS_TOP + 6,
205                &label,
206            );
207        }
208        canvas.text(WIDTH - RIGHT + 8, AXIS_TOP + 6, "time s");
209    }
210
211    /// Encodes the image as an 8-bit RGB PNG.
212    ///
213    /// # Errors
214    ///
215    /// Returns a [`png::EncodingError`] if the encoder fails, which an
216    /// in-memory buffer of the right size never causes.
217    pub fn to_png(&self) -> Result<Vec<u8>, png::EncodingError> {
218        let mut bytes = Vec::new();
219        let width = u32::try_from(WIDTH).unwrap_or(u32::MAX);
220        let height = u32::try_from(HEIGHT).unwrap_or(u32::MAX);
221        let mut encoder = png::Encoder::new(&mut bytes, width, height);
222        encoder.set_color(png::ColorType::Rgb);
223        encoder.set_depth(png::BitDepth::Eight);
224        let mut writer = encoder.write_header()?;
225        writer.write_image_data(&self.pixels())?;
226        writer.finish()?;
227        Ok(bytes)
228    }
229}
230
231/// The first sample and the length of the passage.
232fn passage(samples: &[f64], rate: SampleRate) -> (usize, usize) {
233    let window = cast::round_to_usize(PASSAGE_S.get() * rate.hz().get()).unwrap_or(usize::MAX);
234    if samples.len() <= window {
235        return (0, samples.len());
236    }
237    let params = StftParams::fixed(PANELS[0].frame_len);
238    let hop = params.hop();
239    let active = ActiveFrames::of(samples, params);
240    // prefix[m] counts the active frames before frame m.
241    let mut prefix = vec![0_usize; active.len() + 1];
242    for frame in 0..active.len() {
243        prefix[frame + 1] = prefix[frame] + usize::from(active.is_active(frame));
244    }
245    let span = window / hop;
246    let last_start = (samples.len() - window) / hop;
247    let (mut best, mut best_count) = (0, 0);
248    for start in 0..=last_start {
249        let end = (start + span).min(active.len());
250        let count = prefix[end] - prefix[start.min(end)];
251        if count > best_count {
252            (best, best_count) = (start, count);
253        }
254    }
255    (best * hop, window)
256}
257
258/// The dB matrix of one panel over `len` samples from `first`.
259fn panel_db(
260    samples: &[f64],
261    rate: SampleRate,
262    spec: PanelSpec,
263    first: usize,
264    len: usize,
265) -> Vec<f64> {
266    let params = StftParams::fixed(spec.frame_len);
267    let window_sum: f64 = Window::Hann.samples(spec.frame_len).iter().sum();
268    let bin_hz = params.bin_width(rate).get();
269    let rows = row_bins(spec, bin_hz, params.bins());
270    let mut stft = Stft::new(params);
271    let mut peak = 0.0_f64;
272    for frame in 0..params.frame_count(samples.len()) {
273        // DC is left out, so an offset does not set the scale.
274        for value in stft.frame(samples, frame).iter().skip(1) {
275            peak = peak.max(value.norm() / window_sum);
276        }
277    }
278    // One frame per column, centred on the column's time.
279    let mut columns = Stft::new(params.with_hop(1).unwrap_or(params));
280    let mut db = vec![-RANGE_DB.get(); PLOT_WIDTH * PLOT_HEIGHT];
281    for column in 0..PLOT_WIDTH {
282        let centre = first
283            + cast::floor_to_usize(
284                (cast::usize_to_f64(column) + 0.5) * cast::usize_to_f64(len)
285                    / cast::usize_to_f64(PLOT_WIDTH),
286            )
287            .unwrap_or(0);
288        // With a hop of 1, frame m is centred on sample m.
289        let magnitudes: Vec<f64> = columns
290            .frame(samples, centre)
291            .iter()
292            .map(|value| value.norm() / window_sum)
293            .collect();
294        for (row, bins) in rows.iter().enumerate() {
295            db[row * PLOT_WIDTH + column] = level_db(bins.magnitude(&magnitudes), peak);
296        }
297    }
298    db
299}
300
301/// A level in dB relative to `peak`, clamped to `-RANGE_DB..=0`; silence
302/// and a silent stem read `-RANGE_DB`.
303fn level_db(magnitude: f64, peak: f64) -> f64 {
304    let floor = -RANGE_DB.get();
305    if magnitude > 0.0 && peak > 0.0 {
306        (20.0 * (magnitude / peak).log10()).clamp(floor, 0.0)
307    } else {
308        floor
309    }
310}
311
312/// How a row reads the spectrum.
313#[derive(Clone, Copy, Debug)]
314enum RowBins {
315    /// The largest magnitude of the bins whose centres lie in the row.
316    Max(usize, usize),
317    /// A row narrower than a bin: the magnitude interpolated linearly at
318    /// the row's centre frequency, between bin `k` and `k + 1`.
319    Between(usize, f64),
320    /// Above the Nyquist frequency.
321    Empty,
322}
323
324impl RowBins {
325    fn magnitude(self, magnitudes: &[f64]) -> f64 {
326        match self {
327            Self::Max(first, last) => magnitudes
328                .get(first..=last)
329                .unwrap_or(&[])
330                .iter()
331                .copied()
332                .fold(0.0, f64::max),
333            Self::Between(k, fraction) => match (magnitudes.get(k), magnitudes.get(k + 1)) {
334                (Some(a), Some(b)) => a + fraction * (b - a),
335                _ => 0.0,
336            },
337            Self::Empty => 0.0,
338        }
339    }
340}
341
342/// The bins of each row, top row first. Row `r` covers the frequencies
343/// from `lo (hi / lo)^(1 - (r + 1) / H)` to `lo (hi / lo)^(1 - r / H)`.
344fn row_bins(spec: PanelSpec, bin_hz: f64, bins: usize) -> Vec<RowBins> {
345    let height = cast::usize_to_f64(PLOT_HEIGHT);
346    let edge = |position: f64| {
347        spec.lo.get() * (spec.hi.get() / spec.lo.get()).powf(1.0 - position / height)
348    };
349    (0..PLOT_HEIGHT)
350        .map(|row| {
351            let r = cast::usize_to_f64(row);
352            let (low, high) = (edge(r + 1.0), edge(r));
353            let first = cast::ceil_to_usize(low / bin_hz).unwrap_or(0);
354            let last = cast::ceil_to_usize(high / bin_hz)
355                .unwrap_or(0)
356                .saturating_sub(1);
357            if first >= bins {
358                RowBins::Empty
359            } else if first <= last {
360                RowBins::Max(first, last.min(bins - 1))
361            } else {
362                let centre = (low * high).sqrt() / bin_hz;
363                let k = cast::floor_to_usize(centre).unwrap_or(0);
364                RowBins::Between(k, centre - cast::usize_to_f64(k))
365            }
366        })
367        .collect()
368}
369
370/// The row of a frequency, clamped to the panel.
371fn frequency_row(spec: PanelSpec, frequency: f64) -> usize {
372    let position = (frequency / spec.lo.get()).ln() / (spec.hi.get() / spec.lo.get()).ln();
373    cast::round_to_usize((1.0 - position) * cast::usize_to_f64(PLOT_HEIGHT))
374        .unwrap_or(0)
375        .min(PLOT_HEIGHT - 1)
376}
377
378fn frequency_ticks(spec: PanelSpec) -> &'static [f64] {
379    if spec.hi.get() > 1000.0 {
380        &[
381            20.0, 50.0, 100.0, 200.0, 500.0, 1000.0, 2000.0, 5000.0, 10_000.0, 20_000.0,
382        ]
383    } else {
384        &[20.0, 30.0, 50.0, 100.0, 200.0, 300.0, 500.0]
385    }
386}
387
388fn hz_label(frequency: f64) -> String {
389    if frequency >= 1000.0 {
390        format!("{}k", frequency / 1000.0)
391    } else {
392        format!("{frequency}")
393    }
394}
395
396/// The viridis colour of a level from `-RANGE_DB` (dark) to 0 (bright).
397fn colour(db: f64) -> [u8; 3] {
398    let colour = colorous::VIRIDIS.eval_continuous((db / RANGE_DB.get() + 1.0).clamp(0.0, 1.0));
399    [colour.r, colour.g, colour.b]
400}
401
402/// The colour bar right of the panels with its dB labels.
403fn colour_bar(canvas: &mut Canvas) {
404    let x = WIDTH - RIGHT + 8;
405    let (top, bottom) = (TITLE, AXIS_TOP - 1);
406    let height = bottom - top;
407    for y in top..=bottom {
408        let db = -RANGE_DB.get() * cast::usize_to_f64(y - top) / cast::usize_to_f64(height);
409        for dx in 0..16 {
410            canvas.set(x + dx, y, colour(db));
411        }
412    }
413    canvas.outline(x - 1, top - 1, 18, height + 3);
414    canvas.text(x + 20, top, "0 dB");
415    canvas.text(x + 20, top + height / 2 - 4, "-40");
416    canvas.text(x + 20, bottom - 7, "-80");
417}
418
419struct Canvas {
420    pixels: Vec<u8>,
421}
422
423impl Canvas {
424    fn new() -> Self {
425        Self {
426            pixels: BACKGROUND.repeat(WIDTH * HEIGHT),
427        }
428    }
429
430    fn set(&mut self, x: usize, y: usize, colour: [u8; 3]) {
431        if x < WIDTH && y < HEIGHT {
432            let index = 3 * (y * WIDTH + x);
433            self.pixels[index..index + 3].copy_from_slice(&colour);
434        }
435    }
436
437    fn hline(&mut self, x: usize, y: usize, len: usize) {
438        for dx in 0..len {
439            self.set(x + dx, y, INK);
440        }
441    }
442
443    fn vline(&mut self, x: usize, y: usize, len: usize) {
444        for dy in 0..len {
445            self.set(x, y + dy, INK);
446        }
447    }
448
449    /// A one-pixel frame around a `width` by `height` box at `(x, y)`.
450    fn outline(&mut self, x: usize, y: usize, width: usize, height: usize) {
451        for dx in 0..width {
452            self.set(x + dx, y, FRAME);
453            self.set(x + dx, y + height - 1, FRAME);
454        }
455        for dy in 0..height {
456            self.set(x, y + dy, FRAME);
457            self.set(x + width - 1, y + dy, FRAME);
458        }
459    }
460
461    /// Text in the 8x8 font with its top left corner at `(x, y)`.
462    fn text(&mut self, x: usize, y: usize, text: &str) {
463        for (index, character) in text.chars().enumerate() {
464            let glyph = font8x8::BASIC_FONTS.get(character).unwrap_or([0; 8]);
465            for (dy, bits) in glyph.iter().enumerate() {
466                for dx in 0..8 {
467                    if bits >> dx & 1 == 1 {
468                        self.set(x + 8 * index + dx, y + dy, INK);
469                    }
470                }
471            }
472        }
473    }
474}