scripts/select_sequences.py

"""Survey val sequences with the released model to pick demo candidates.

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

    cd ../pyronear/temporal-model/train
    uv run python <this file> [--sample 12] [--seed 7]

Prints one line per sequence so we can pick a diverse set for the space:
clear/faint positives, positives with interpolated gaps, and false-positive
look-alikes where the detector fires but the temporal model says no.
"""

import argparse
import random
from pathlib import Path

from huggingface_hub import hf_hub_download
from temporal_model.core.model import BboxTubeTemporalModel

MODEL_VERSION = "0.2.0"
VAL_ROOT = Path(__file__).resolve().parents[3] / (
    "pyronear/temporal-model/train/data/01_raw/datasets/val"
)


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 survey(model: BboxTubeTemporalModel, seq_dir: Path, label: str) -> str:
    frames = sorted((seq_dir / "images").glob("*.jpg"))
    out = model.predict_sequence(frames)
    det = out.details
    kept = det["tubes"]["kept"]
    best_prob = max((t["probability"] or 0.0) for t in kept) if kept else 0.0
    gaps = sum(e["is_gap"] for t in kept for e in t["entries"])
    return (
        f"{label:8s} {seq_dir.name:55s} "
        f"verdict={'SMOKE' if out.is_positive else 'no   '} "
        f"p={best_prob:.3f} kept={len(kept)} "
        f"cand={det['tubes']['num_candidates']:2d} "
        f"trigger={out.trigger_frame_index} gaps={gaps} "
        f"frames={len(frames)}"
    )


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--sample", type=int, default=12)
    parser.add_argument("--seed", type=int, default=7)
    args = parser.parse_args()

    rng = random.Random(args.seed)
    model = load_model()
    for label in ("wildfire", "fp"):
        dirs = sorted((VAL_ROOT / label).iterdir())
        for seq_dir in rng.sample(dirs, min(args.sample, len(dirs))):
            try:
                print(survey(model, seq_dir, label), flush=True)
            except Exception as e:  # noqa: BLE001 - survey keeps going
                print(f"{label:8s} {seq_dir.name:55s} ERROR {e}", flush=True)


if __name__ == "__main__":
    main()