Live heatmap display through debug rtt

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<ambient>-?\d+(?:\.\d+)?).*?"
+    r"min_c=(?P<min>-?\d+(?:\.\d+)?).*?"
+    r"max_c=(?P<max>-?\d+(?:\.\d+)?).*?"
+    r"pixels=\[(?P<pixels>[^\]]*)\]"
+)
+
+
+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()

src/main.rs

@@ -53,13 +53,20 @@ async fn main(_spawner: Spawner) {
 
     sensor::wake_up_and_load_trim(&mut i2c, &cal).await;
 
+    // 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!(
-    info!("ambient temperature: {} C", heatmap.ambient_c);
+            "frame ambient_c={} min_c={} max_c={} pixels={}",
-    info!("heatmap range: {}..{} C", heatmap.min_c, heatmap.max_c);
+            heatmap.ambient_c,
-    for row in 0..calibration::GRID {
+            heatmap.min_c,
-        info!("row {}: {}", row, heatmap.temperature_c[row]);
+            heatmap.max_c,
+            heatmap.flatten()
+        );
     }
 }
 

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<const ROWS: usize, const COLS: usize>(
     cal: &Calibration,