From c0eacb9103a82c2ced0efa7de9b34ccd673869f6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Piotr=20Stefa=C5=84ski?= Date: Sun, 5 Jul 2026 22:32:57 +0200 Subject: [PATCH] Live heatmap display through debug rtt --- live_heatmap.py | 185 ++++++++++++++++++++++++++++++++++++++++++++++++ src/main.rs | 21 ++++-- src/thermal.rs | 14 ++++ 3 files changed, 213 insertions(+), 7 deletions(-) create mode 100644 live_heatmap.py diff --git a/live_heatmap.py b/live_heatmap.py new file mode 100644 index 0000000..9381cd7 --- /dev/null +++ b/live_heatmap.py @@ -0,0 +1,185 @@ +#!/usr/bin/env python3 +"""Live 16x16 thermal heatmap viewer. + +Reads the firmware's RTT/defmt log stream and updates a heatmap window in +real time, one frame at a time. The firmware (see `src/main.rs`) prints one +line per frame that looks like: + + 12.345 [INFO ] frame ambient_c=23.5 min_c=10.2 max_c=45.6 pixels=[23.1, ...] (tts_test ...) + +Everything else on the line (timestamp, log level, source location) and any +other log lines (warnings, startup messages, ...) are ignored. + +Usage: + # live, straight from the device (uses the runner in .cargo/config.toml, + # i.e. `probe-rs run --chip STM32U585CI`): + cargo run --release | python3 live_heatmap.py + + # replay a previously captured log file instead of a live device: + python3 live_heatmap.py --replay output.txt +""" + +import argparse +import re +import select +import sys + +import matplotlib.pyplot as plt +import numpy as np + +GRID = 16 + +# Fixed color-scale range (\u00b0C). Kept constant across frames (rather than +# autoscaling to each frame's min/max) so colors stay comparable over time; +# override with --vmin/--vmax if your scene runs outside this range. +DEFAULT_VMIN_C = 0.0 +DEFAULT_VMAX_C = 100.0 + +_ANSI_RE = re.compile(r"\x1b\[[0-9;]*m") +_FRAME_RE = re.compile( + r"ambient_c=(?P-?\d+(?:\.\d+)?).*?" + r"min_c=(?P-?\d+(?:\.\d+)?).*?" + r"max_c=(?P-?\d+(?:\.\d+)?).*?" + r"pixels=\[(?P[^\]]*)\]" +) + + +def parse_frame(line): + """Return (ambient, min, max, 16x16 array), or None if `line` isn't a frame line.""" + line = _ANSI_RE.sub("", line) + match = _FRAME_RE.search(line) + if not match: + return None + + pixels = [float(v) for v in match.group("pixels").split(",") if v.strip()] + if len(pixels) != GRID * GRID: + return None + + grid = np.array(pixels, dtype=float).reshape(GRID, GRID) + ambient = float(match.group("ambient")) + vmin = float(match.group("min")) + vmax = float(match.group("max")) + return ambient, vmin, vmax, grid + + +def read_latest_frame(stream): + """Block until at least one frame line is available on `stream`, then + drain any further lines that are *already* buffered and return only the + most recently parsed frame. + + This is what keeps the viewer live: if frames arrive faster than they + can be drawn, we skip straight to the newest one each time instead of + rendering every buffered frame in order and steadily falling behind. + + Returns `None` once `stream` reaches EOF. + """ + frame = None + while True: + line = stream.readline() + if line == "": + return frame # EOF + parsed = parse_frame(line) + if parsed is not None: + frame = parsed + # non-blocking peek: is there already more data waiting? + more_data_ready, _, _ = select.select([stream], [], [], 0) + if not more_data_ready and frame is not None: + return frame + # otherwise keep draining (or, if nothing parsed yet, block again + # on the next readline() until a new line arrives) + + +class LiveHeatmap: + """A matplotlib window that redraws itself each time `update` is called. + + The color scale (`vmin`/`vmax`) is fixed for the lifetime of the window, + so colors stay comparable from frame to frame; only the pixel data and + the title (which still reports each frame's actual min/max) change. + """ + + def __init__(self, cmap, vmin, vmax, annotate=True): + plt.ion() + self.fig, ax = plt.subplots(figsize=(8, 7)) + self.image = ax.imshow( + np.zeros((GRID, GRID)), cmap=cmap, interpolation="nearest", vmin=vmin, vmax=vmax + ) + self.colorbar = self.fig.colorbar(self.image, ax=ax, label="\u00b0C") + self.title = ax.set_title("waiting for first frame...") + ax.set_xticks(range(GRID)) + ax.set_yticks(range(GRID)) + + # midpoint of the fixed color scale, used to pick readable text color + self.mid = (vmin + vmax) / 2 + self.annotations = None + if annotate: + self.annotations = [ + [ + ax.text(col, row, "", ha="center", va="center", fontsize=6, color="white") + for col in range(GRID) + ] + for row in range(GRID) + ] + + self.fig.tight_layout() + self.fig.canvas.draw() + self.fig.canvas.flush_events() + self.frame_count = 0 + + def update(self, ambient, frame_min, frame_max, grid): + self.frame_count += 1 + self.image.set_data(grid) + self.title.set_text( + f"frame {self.frame_count} ambient {ambient:.1f}\u00b0C " + f"range {frame_min:.1f}..{frame_max:.1f}\u00b0C" + ) + if self.annotations is not None: + for row in range(GRID): + for col in range(GRID): + value = grid[row, col] + text = self.annotations[row][col] + text.set_text(f"{value:.0f}") + text.set_color("white" if value < self.mid else "black") + self.fig.canvas.draw_idle() + self.fig.canvas.flush_events() + + def is_open(self): + return plt.fignum_exists(self.fig.number) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--replay", help="replay a captured log file instead of reading stdin") + parser.add_argument("--replay-fps", type=float, default=5.0, help="frame rate when replaying a file") + parser.add_argument("--cmap", default="inferno", help="matplotlib colormap") + parser.add_argument("--vmin", type=float, default=DEFAULT_VMIN_C, help="fixed color-scale minimum (\u00b0C)") + parser.add_argument("--vmax", type=float, default=DEFAULT_VMAX_C, help="fixed color-scale maximum (\u00b0C)") + parser.add_argument("--no-annot", action="store_true", help="do not print the value inside each cell") + args = parser.parse_args() + + heatmap = LiveHeatmap(args.cmap, args.vmin, args.vmax, annotate=not args.no_annot) + try: + if args.replay: + # deliberately paced, frame by frame, at --replay-fps + with open(args.replay) as source: + for line in source: + if not heatmap.is_open(): + break + frame = parse_frame(line) + if frame is None: + continue + heatmap.update(*frame) + plt.pause(1.0 / args.replay_fps) + else: + # live: always show the newest frame, never a backlog + while heatmap.is_open(): + frame = read_latest_frame(sys.stdin) + if frame is None: + break # EOF: the device stream ended + heatmap.update(*frame) + plt.pause(0.001) # let the GUI event loop breathe + except KeyboardInterrupt: + pass + + +if __name__ == "__main__": + main() diff --git a/src/main.rs b/src/main.rs index 1a9c947..a8f21af 100644 --- a/src/main.rs +++ b/src/main.rs @@ -53,13 +53,20 @@ async fn main(_spawner: Spawner) { sensor::wake_up_and_load_trim(&mut i2c, &cal).await; - let frame = sensor::Frame::capture(&mut i2c).await; - let heatmap = thermal::process_frame(&cal, &frame, &table_tn114::TABLE); - - info!("ambient temperature: {} C", heatmap.ambient_c); - info!("heatmap range: {}..{} C", heatmap.min_c, heatmap.max_c); - for row in 0..calibration::GRID { - info!("row {}: {}", row, heatmap.temperature_c[row]); + // Stream one frame per line, forever, so a host-side tool can render it + // live from the RTT log (see `live_heatmap.py`). Everything needed to + // draw the heatmap is packed into this single `info!` call: the whole + // 16x16 grid is flattened into one array so it survives as one line. + loop { + let frame = sensor::Frame::capture(&mut i2c).await; + let heatmap = thermal::process_frame(&cal, &frame, &table_tn114::TABLE); + info!( + "frame ambient_c={} min_c={} max_c={} pixels={}", + heatmap.ambient_c, + heatmap.min_c, + heatmap.max_c, + heatmap.flatten() + ); } } diff --git a/src/thermal.rs b/src/thermal.rs index 4aacc37..70da1e0 100644 --- a/src/thermal.rs +++ b/src/thermal.rs @@ -20,6 +20,20 @@ pub struct Heatmap { pub temperature_c: [[f32; GRID]; GRID], } +impl Heatmap { + /// Flatten the grid into one row-major array, so it can be sent as a + /// single log line for real-time streaming (see `main.rs`). + pub fn flatten(&self) -> [f32; GRID * GRID] { + let mut flat = [0f32; GRID * GRID]; + for row in 0..GRID { + for col in 0..GRID { + flat[row * GRID + col] = self.temperature_c[row][col]; + } + } + flat + } +} + /// Compensate one raw frame into a full object-temperature heatmap. pub fn process_frame( cal: &Calibration,