scripts/precompute.py

"""Precompute demo payloads for the curated sequences.

Copies (resized) frames into data/sequences/<seq-id>/, runs the released
temporal model on those exact copies (so the shipped results match what the
space's live run sees), and writes per-sequence payloads under
data/precomputed/<seq-id>/:

  - result.json   verdict, trigger frame, full BboxTubeDetails payload
  - filmstrip_tube<id>.jpg   stabilized patch strip per kept tube

Run from the temporal-model train package (its venv has core + torch):

    cd ../pyronear/temporal-model/train
    uv run python <this file>
"""

import json
import sys
from pathlib import Path

from huggingface_hub import hf_hub_download
from PIL import Image, ImageDraw
from temporal_model.core.crop import crop_and_resize, expand_bbox, norm_bbox_to_pixel_square
from temporal_model.core.model import BboxTubeTemporalModel

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
import render  # noqa: E402 - shared overlay/animation rendering

MODEL_VERSION = "0.2.0"
MAX_FRAME_WIDTH = 1280
JPEG_QUALITY = 87
PATCH = 224

HERE = Path(__file__).resolve().parent
SPACE = HERE.parent
DATA = SPACE / "data"
VAL_ROOT = SPACE.parents[1] / "pyronear/temporal-model/train/data/01_raw/datasets/val"

# (label, sequence id) — picked with select_sequences.py for narrative variety:
# a clean growing plume, a fragmented/gappy plume, a look-alike that builds a
# tube but is rejected by the classifier, and detector noise that never
# survives tube building.
SEQUENCES = [
    ("wildfire", "pyronear-sdis-07_marguerite_282_2024-02-23T12-38-40"),
    ("wildfire", "pyronear-sdis-07_brison_200_2024-02-21T15-16-17"),
    ("fp", "pyronear-force-06_cabanelle_291_2023-10-15T11-06-20"),
    ("fp", "pyronear-sdis-07_brison_110_2024-01-16T11-00-04"),
]


def load_model() -> BboxTubeTemporalModel:
    zip_path = hf_hub_download(
        repo_id="pyronear/temporal-model",
        filename="model.zip",
        revision=f"v{MODEL_VERSION}",
    )
    return BboxTubeTemporalModel.from_package(Path(zip_path))


def copy_frames(src_dir: Path, dst_dir: Path) -> list[Path]:
    dst_dir.mkdir(parents=True, exist_ok=True)
    out = []
    for src in sorted(src_dir.glob("*.jpg")):
        dst = dst_dir / src.name
        img = Image.open(src).convert("RGB")
        if img.width > MAX_FRAME_WIDTH:
            img = img.resize(
                (MAX_FRAME_WIDTH, round(img.height * MAX_FRAME_WIDTH / img.width)),
                Image.LANCZOS,
            )
        img.save(dst, quality=JPEG_QUALITY)
        out.append(dst)
    return out


PER_ROW = 7


def render_filmstrip(
    tube: dict, frame_paths: list[Path], context_factor: float, threshold: float
) -> Image.Image:
    """Stabilized patch per (non-padded) tube entry — the classifier's view.

    Wraps at PER_ROW patches per row so the image keeps a screen-friendly
    aspect ratio in the app. Border color reflects the tube's verdict
    (orange = judged smoke, gray = rejected).
    """
    window = tube["stabilized_window"]
    entries = [e for e in tube["entries"] if e["frame_idx"] < len(frame_paths)]
    gap = 6
    label_h = 22
    cols = min(PER_ROW, len(entries))
    rows = -(-len(entries) // PER_ROW)
    cell_h = PATCH + label_h
    strip = Image.new(
        "RGB",
        (cols * (PATCH + gap) - gap, rows * (cell_h + gap) - gap),
        "white",
    )
    draw = ImageDraw.Draw(strip)
    font = render.font(15)
    color = render.tube_color(tube, threshold)
    import numpy as np

    for i, entry in enumerate(entries):
        img = np.array(Image.open(frame_paths[entry["frame_idx"]]).convert("RGB"))
        h, w, _ = img.shape
        cx, cy, bw, bh = expand_bbox(*window, context_factor)
        box = norm_bbox_to_pixel_square(cx, cy, bw, bh, w, h)
        patch = Image.fromarray(crop_and_resize(img, box, PATCH))
        x = (i % PER_ROW) * (PATCH + gap)
        y = (i // PER_ROW) * (cell_h + gap)
        strip.paste(patch, (x, y + label_h))
        draw.rectangle([x, y + label_h, x + PATCH - 1, y + label_h + PATCH - 1], outline=color, width=4)
        suffix = " (interp.)" if entry["is_gap"] else ""
        draw.text((x + 2, y + 2), f"frame {entry['frame_idx'] + 1}{suffix}", fill=render.INK, font=font)
    return strip


def main() -> None:
    model = load_model()
    context_factor = model._cfg["model_input"]["context_factor"]
    index = []
    for label, seq_id in SEQUENCES:
        print(f"=== {seq_id}", flush=True)
        frame_paths = copy_frames(VAL_ROOT / label / seq_id / "images", DATA / "sequences" / seq_id)
        out = model.predict(model.load_sequence(frame_paths), compute_trigger=True)

        out_dir = DATA / "precomputed" / seq_id
        out_dir.mkdir(parents=True, exist_ok=True)
        kept = out.details["tubes"]["kept"]
        threshold = out.details["decision"]["threshold"]
        for tube in kept:
            if tube["stabilized_window"] is None:
                continue
            strip = render_filmstrip(tube, frame_paths, context_factor, threshold)
            strip.save(out_dir / f"filmstrip_tube{tube['tube_id']}.jpg", quality=90)

        payload = {
            "sequence_id": seq_id,
            "ground_truth": label,
            "model_version": MODEL_VERSION,
            "frames": [p.name for p in frame_paths],
            "is_positive": out.is_positive,
            "trigger_frame_index": out.trigger_frame_index,
            "details": out.details,
        }
        (out_dir / "result.json").write_text(json.dumps(payload, indent=2))
        render.save_animation(payload, frame_paths, out_dir / "animation.webp")
        index.append(
            {
                "sequence_id": seq_id,
                "ground_truth": label,
                "is_positive": out.is_positive,
                "num_frames": len(frame_paths),
            }
        )
        print(
            f"    verdict={'SMOKE' if out.is_positive else 'no smoke'} "
            f"trigger={out.trigger_frame_index} kept={len(kept)}",
            flush=True,
        )
    (DATA / "precomputed" / "index.json").write_text(json.dumps(index, indent=2))


if __name__ == "__main__":
    main()