ppak10/Agentic-SLS-Telemetry

1#!/usr/bin/env python32"""Render preview MP4s for each build's parquet.3 4Reads data/ticks/{build_id:03d}.parquet and writes:5  previews/{build_id:03d}/chamber.mp4    (optical)6  previews/{build_id:03d}/thermal.mp47  previews/{build_id:03d}/galvo.mp48  previews/{build_id:03d}/composite.mp4  (1x3 panel: chamber | thermal | galvo)9 10Playback is at 10 fps (= tick rate), so the video duration matches the build's11wall-clock duration. Null frames are forward-filled with the most recent12captured frame of that kind, so the preview keeps moving even during the13sparse stretches (heating, idle).14 15Source is the dataset's own parquet — no recorder repo needed. The trade-off16is that ~half of upstream-captured frames don't survive the per-tick attach,17so motion looks chunkier than playing the raw recorder frames would.18 19The thermal panel is rendered from the raw `bedmatrix` IR grid (inferno colormap20over a fixed absolute °C range; see _thermal.py) whenever it's present — builds21013+. The three earliest builds (001/002/012) predate the bedmatrix stream and22fall back to the legacy pre-rendered `frame_thermal` GIF.23 24Usage:25    uv run scripts/previews/01_render.py            # all builds in data/ticks/26    uv run scripts/previews/01_render.py 13 26      # specific build ids27"""28import io29import os30import sys31from pathlib import Path32 33# Point imageio-ffmpeg at the system ffmpeg (Ubuntu's build includes h264_nvenc;34# the bundled imageio-ffmpeg binary doesn't). Must be set BEFORE importing imageio.35os.environ.setdefault("IMAGEIO_FFMPEG_EXE", "/usr/bin/ffmpeg")36 37import imageio.v2 as iio38import numpy as np39import pyarrow.parquet as pq40from PIL import Image41 42sys.path.insert(0, str(Path(__file__).parent.parent))43sys.path.insert(0, str(Path(__file__).parent))44from _lib import DATA_DIR45from _thermal import bedmatrix_to_image46 47OUTPUT_DIR = DATA_DIR.parent / "previews"48TICKS_DIR = DATA_DIR / "ticks"49# Order here = left-to-right order in the composite panel.50FRAME_KINDS = ("chamber", "thermal", "galvo")51# The thermal panel is rendered from the raw `bedmatrix` IR grid when present52# (builds 013+), falling back to the legacy `frame_thermal` GIF for the three53# earliest builds (001/002/012) that predate the bedmatrix stream.54# Per-kind rotation applied post-decode (degrees clockwise). The chamber camera55# is mounted sideways on the printer, so its raw frames need a 90° CW correction56# before rendering. Other kinds are captured already in display orientation.57KIND_ROTATION_CW = {"chamber": 90}58FPS = 1059PANEL_HEIGHT = 480  # all panels resized to this height; widths float to preserve aspect ratio.60BATCH_ROWS = 1000   # pyarrow row-group iter chunk; bounds peak memory per build.61 62 63def _even(n: int) -> int:64    """H.264 requires even dimensions; round up if odd."""65    return n if n % 2 == 0 else n + 166 67 68def _decode_raw(frame_struct, kind: str) -> Image.Image | None:69    """Decode an HF Image struct to a PIL Image (RGB), applying KIND_ROTATION_CW70    if the kind needs orientation correction. None when missing."""71    if frame_struct is None:72        return None73    b = frame_struct.get("bytes")74    if b is None:75        return None76    try:77        img = Image.open(io.BytesIO(b)).convert("RGB")78    except Exception:79        return None80    rot = KIND_ROTATION_CW.get(kind, 0)81    # Image.Transpose.ROTATE_N rotates N degrees CCW, so CW=N → ROTATE_(360-N).82    if rot == 90:83        img = img.transpose(Image.Transpose.ROTATE_270)84    elif rot == 180:85        img = img.transpose(Image.Transpose.ROTATE_180)86    elif rot == 270:87        img = img.transpose(Image.Transpose.ROTATE_90)88    return img89 90 91def _decode_cell(cell, kind: str) -> Image.Image | None:92    """Decode one struct cell to a PIL RGB image, dispatching on struct shape:93    a `bedmatrix` struct (has 'values') renders as an inferno heatmap; a frame94    Image struct (has 'bytes') decodes + orientation-corrects. None when missing."""95    if cell is None:96        return None97    if "values" in cell:98        return bedmatrix_to_image(cell)99    return _decode_raw(cell, kind)100 101 102def _thermal_column(parquet_path: Path) -> str:103    """Which column feeds the thermal panel for this build: 'bedmatrix' when the104    raw IR matrix has any non-null cell, else the legacy 'frame_thermal'. The105    bedmatrix stream started at build 013, so 001/002/012 fall back to the GIF."""106    pf = pq.ParquetFile(parquet_path)107    if "bedmatrix" not in {f.name for f in pf.schema_arrow}:108        return "frame_thermal"109    for batch in pf.iter_batches(columns=["bedmatrix"], batch_size=BATCH_ROWS):110        for cell in batch.column("bedmatrix").to_pylist():111            if cell is not None:112                return "bedmatrix"113    return "frame_thermal"114 115 116def _columns_for_build(parquet_path: Path) -> dict[str, str]:117    """Map each panel kind → the parquet column that feeds it for this build."""118    return {119        "chamber": "frame_chamber",120        "thermal": _thermal_column(parquet_path),121        "galvo":   "frame_galvo",122    }123 124 125def _fit_to_canvas(img: Image.Image, canvas_w: int) -> np.ndarray:126    """Letterbox `img` into a (canvas_w × PANEL_HEIGHT) gray canvas, preserving aspect.127    Locking on the canvas dims is required because some builds have frames whose128    source dimensions drift mid-build (e.g. thermal IR config changes)."""129    sw, sh = img.size130    scale = min(canvas_w / sw, PANEL_HEIGHT / sh)131    new_w = max(1, round(sw * scale))132    new_h = max(1, round(sh * scale))133    fitted = img.resize((new_w, new_h), Image.BILINEAR)134    canvas = Image.new("RGB", (canvas_w, PANEL_HEIGHT), (20, 20, 20))135    canvas.paste(fitted, ((canvas_w - new_w) // 2, (PANEL_HEIGHT - new_h) // 2))136    return np.asarray(canvas)137 138 139def _placeholder(width: int) -> np.ndarray:140    """Gray panel shown before the first frame of a kind has been seen."""141    return np.full((PANEL_HEIGHT, width, 3), 20, dtype=np.uint8)142 143 144def _iter_struct_batches(parquet_path: Path, columns: list[str]):145    """Yield (n_rows, dict[col -> list[struct|None]]) per row-group batch."""146    pf = pq.ParquetFile(parquet_path)147    for batch in pf.iter_batches(columns=columns, batch_size=BATCH_ROWS):148        cols = {c: batch.column(c).to_pylist() for c in columns}149        yield batch.num_rows, cols150 151 152def _open_writer(out_path: Path):153    """Open an MP4 writer. Uses NVENC on the K620 GPUs when available, falling154    back to libx264 if the env var RENDER_CPU=1 forces software encode."""155    out_path.parent.mkdir(parents=True, exist_ok=True)156    if os.environ.get("RENDER_CPU") == "1":157        return iio.get_writer(158            out_path, fps=FPS, codec="libx264",159            macro_block_size=1, quality=6,  # ~CRF 23-ish160        )161    # NVENC path. Quality target ~CRF 23 via constant-quality VBR. `p4` is the162    # balanced preset on the new naming; older NVENC firmware may report this163    # as "medium". yuv420p forced because rgb24 input → NVENC needs 4:2:0.164    return iio.get_writer(165        out_path, fps=FPS, codec="h264_nvenc",166        macro_block_size=1,167        ffmpeg_params=[168            "-preset", "p4",169            "-rc", "vbr",170            "-cq", "23",171            "-pix_fmt", "yuv420p",172        ],173    )174 175 176def _probe_first_frame_width(parquet_path: Path, kind: str, col: str) -> int | None:177    """Find the first non-null cell in `col`, compute the locked canvas width178    from its post-decode aspect ratio (height = PANEL_HEIGHT). Returns None if179    no frame of that kind ever appears."""180    for n_rows, cols in _iter_struct_batches(parquet_path, [col]):181        for struct in cols[col]:182            img = _decode_cell(struct, kind)183            if img is not None:184                sw, sh = img.size185                return _even(max(2, round(sw * PANEL_HEIGHT / sh)))186    return None187 188 189def render_per_kind(parquet_path: Path, kind: str, out_path: Path, col: str) -> int:190    """Write one MP4 of a single panel kind, forward-filled. Returns frame count."""191    width = _probe_first_frame_width(parquet_path, kind, col)192    if width is None:193        # No frames of this kind exist for this build; nothing meaningful to render.194        print(f"  {kind:8s}: no frames of this kind, skipping")195        return 0196    last: np.ndarray | None = None197    frames_written = 0198    with _open_writer(out_path) as writer:199        for n_rows, cols in _iter_struct_batches(parquet_path, [col]):200            for struct in cols[col]:201                img = _decode_cell(struct, kind)202                if img is not None:203                    last = _fit_to_canvas(img, width)204                writer.append_data(last if last is not None else _placeholder(width))205                frames_written += 1206    return frames_written207 208 209def render_composite(parquet_path: Path, out_path: Path, cols_map: dict[str, str]) -> int:210    """Write the 1×3 composite. All three panels share the tick timeline."""211    # Lock canvas widths up front so all subsequent frames letterbox into a fixed shape.212    widths = {k: (_probe_first_frame_width(parquet_path, k, cols_map[k]) or PANEL_HEIGHT)213              for k in FRAME_KINDS}214    cols_to_read = [cols_map[k] for k in FRAME_KINDS]215    last: dict[str, np.ndarray | None] = {k: None for k in FRAME_KINDS}216    frames_written = 0217    with _open_writer(out_path) as writer:218        for n_rows, cols in _iter_struct_batches(parquet_path, cols_to_read):219            for i in range(n_rows):220                for k in FRAME_KINDS:221                    img = _decode_cell(cols[cols_map[k]][i], k)222                    if img is not None:223                        last[k] = _fit_to_canvas(img, widths[k])224                panels = [225                    last[k] if last[k] is not None else _placeholder(widths[k])226                    for k in FRAME_KINDS227                ]228                writer.append_data(np.hstack(panels))229                frames_written += 1230    return frames_written231 232 233def process_build(build_id: int, kinds: set[str]) -> None:234    parquet_path = TICKS_DIR / f"{build_id:03d}.parquet"235    if not parquet_path.exists():236        print(f"build {build_id:03d}: no parquet, skipping")237        return238    build_dir = OUTPUT_DIR / f"{build_id:03d}"239    print(f"build {build_id:03d}: rendering → {build_dir.relative_to(Path.cwd())}/")240    cols_map = _columns_for_build(parquet_path)241    if ("thermal" in kinds or "composite" in kinds):242        print(f"  thermal source: {cols_map['thermal']}")243    for kind in FRAME_KINDS:244        if kind not in kinds:245            continue246        out = build_dir / f"{kind}.mp4"247        n = render_per_kind(parquet_path, kind, out, cols_map[kind])248        # render_per_kind skips writing entirely when no frames of this kind exist249        # (and prints its own "skipping" line). Guard stat to avoid FileNotFoundError.250        if out.exists():251            print(f"  {kind:8s}: {n:>7,} frames → {out.name} ({out.stat().st_size:,} bytes)")252    if "composite" in kinds:253        out = build_dir / "composite.mp4"254        n = render_composite(parquet_path, out, cols_map)255        if out.exists():256            print(f"  composite: {n:>7,} frames → {out.name} ({out.stat().st_size:,} bytes)")257 258 259def main():260    import argparse261    parser = argparse.ArgumentParser(description=__doc__,262                                     formatter_class=argparse.RawDescriptionHelpFormatter)263    parser.add_argument("build_ids", nargs="*", type=int,264                        help="Specific build IDs to render (default: all in data/ticks/)")265    parser.add_argument("--kinds", default="chamber,thermal,galvo,composite",266                        help="Comma-separated outputs to render. Default: all four. "267                             "Useful for re-rendering just one kind, e.g. "268                             "--kinds chamber,composite after a chamber-orientation change.")269    args = parser.parse_args()270    kinds = set(args.kinds.split(","))271    unknown = kinds - (set(FRAME_KINDS) | {"composite"})272    if unknown:273        parser.error(f"unknown --kinds values: {sorted(unknown)}; "274                     f"valid: {sorted(set(FRAME_KINDS) | {'composite'})}")275    targets = args.build_ids or sorted(int(p.stem) for p in TICKS_DIR.glob("*.parquet"))276    for bid in targets:277        process_build(bid, kinds)278 279 280if __name__ == "__main__":281    main()282