scripts/build_coverage_index.py

#!/usr/bin/env python
"""Turn the coverage manifest + downloaded mp4s into a browsable bundle.

Emits, under ``--out-dir``:
    README.md          the axes, what is and is not covered, and the full table
    index.csv          one row per episode, for sorting/filtering elsewhere
    manifest.json      the picker's own output, copied verbatim
    thumbs/<uuid>.jpg  a 3-frame strip (5% / 45% / 90% of the clip) per episode
    sheets/<class>.jpg one contact sheet per object-physics class

The per-class sheets exist because the object class is the axis that decides
whether this project's representation applies at all, so it is the one worth
being able to eyeball as a group rather than episode by episode.

Nothing here re-derives labels: everything comes from the manifest written by
``build_droid_coverage_sample.py``, so the bundle can never disagree with the
selection that produced it.

Usage:
    PYTHONPATH=src python scripts/build_coverage_index.py \\
        --manifest outputs/droid_coverage/manifest.json \\
        --video-root data/droid_coverage56 --out-dir outputs/droid_coverage/bundle
"""

from __future__ import annotations

import argparse
import csv
import json
import shutil
from collections import Counter, defaultdict
from pathlib import Path

import cv2
import imageio.v3 as iio
import numpy as np

_FRACS = (0.05, 0.45, 0.90)
_TILE_W = 360


def strip_for(video: Path, label: str) -> tuple[np.ndarray, int, tuple[int, int]]:
    v = np.asarray(iio.imread(video, plugin="pyav"))
    n = len(v)
    h = int(round(_TILE_W * v.shape[1] / v.shape[2]))
    tiles = [cv2.resize(v[min(n - 1, int(n * f))], (_TILE_W, h)) for f in _FRACS]
    row = np.concatenate(tiles, axis=1)
    cv2.rectangle(row, (0, 0), (len(label) * 9 + 8, 20), (0, 0, 0), -1)
    cv2.putText(row, label, (4, 14), cv2.FONT_HERSHEY_SIMPLEX, 0.42, (255, 255, 255), 1, cv2.LINE_AA)
    return row, n, (v.shape[2], v.shape[1])


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--manifest", required=True, type=Path)
    ap.add_argument("--video-root", required=True, type=Path)
    ap.add_argument("--out-dir", required=True, type=Path)
    args = ap.parse_args()

    rows = json.loads(args.manifest.read_text())
    out = args.out_dir
    (out / "thumbs").mkdir(parents=True, exist_ok=True)
    (out / "sheets").mkdir(parents=True, exist_ok=True)
    shutil.copy(args.manifest, out / "manifest.json")

    by_class: dict[str, list[np.ndarray]] = defaultdict(list)
    table: list[dict] = []
    missing: list[str] = []

    for r in rows:
        vids = list((args.video_root / r["uuid"] / "recordings" / "MP4").glob("*.mp4"))
        if not vids:
            missing.append(r["uuid"])
            continue
        label = f"{r['object_class']} | {r['task_class']} | {r['lab']}" + ("" if r["success"] else " | FAIL")
        row, n, (w, h) = strip_for(vids[0], label)
        cv2.imwrite(str(out / "thumbs" / f"{r['uuid']}.jpg"), row[..., ::-1],
                    [cv2.IMWRITE_JPEG_QUALITY, 88])
        by_class[r["object_class"]].append(row)
        table.append({
            "uuid": r["uuid"], "object_class": r["object_class"], "task_class": r["task_class"],
            "lab": r["lab"], "scene_id": r["scene_id"], "success": r["success"],
            "frames": n, "resolution": f"{w}x{h}",
            "current_task": r["current_task"].replace("\n", " ")[:160],
            "instruction_1": (r["instructions"][0] if r["instructions"] else "").replace("\n", " ")[:160],
        })

    for cls, strips in by_class.items():
        cv2.imwrite(str(out / "sheets" / f"{cls}.jpg"),
                    np.concatenate(strips, axis=0)[..., ::-1], [cv2.IMWRITE_JPEG_QUALITY, 85])

    with (out / "index.csv").open("w", newline="") as f:
        wcsv = csv.DictWriter(f, fieldnames=list(table[0].keys()))
        wcsv.writeheader()
        wcsv.writerows(table)

    cnt = {k: Counter(t[k] for t in table) for k in ("object_class", "task_class", "lab", "success")}
    lens = sorted(t["frames"] for t in table)

    def block(name: str) -> str:
        return "\n".join(f"| {k} | {v} |" for k, v in sorted(cnt[name].items(), key=lambda x: str(x[0])))

    readme = f"""# DROID coverage sample ({len(table)} episodes)

Chosen to **cover DROID's axes of variation**, not sampled uniformly. A uniform draw
returns mostly TRI-lab pick-and-place of rigid objects; this spreads over the
properties that change what a geometry pipeline has to do.

Only camera `ext1` is included. Every episode here is also covered by PointWorld's
camera set, so annotations line up if you go further.

## Two properties of DROID worth knowing before you look

1. **`current_task` is usually a category template, not a description.** Strings like
   `"Move object into or out of container (ex: drawer, clothes hamper, plate, trashcan, washer)"`
   or `"Do any task, and then reset the scene."` are the norm. Object identity has to
   come from the separate human-annotation file
   (`aggregated-annotations-030724.json`, 3 instructions per episode).
2. **The human annotations cover successful episodes only.** Measured over 1,153
   episodes' metadata: all 832 with annotations are successes, all 125 failures have
   none. So any object-based filter silently selects for success. The
   `unknown_no_annotation` rows below are failure episodes drawn deliberately, and
   their object identity is genuinely unknown.

## Object physics class

The axis that decides whether one rigid SE(3) pose per frame is even the right model.

| class | n |
|---|---|
{block("object_class")}

## Task template

| class | n |
|---|---|
{block("task_class")}

## Lab

| lab | n |
|---|---|
{block("lab")}

## Outcome

| success | n |
|---|---|
{block("success")}

## Episode length (frames)

min {lens[0]} / median {lens[len(lens)//2]} / max {lens[-1]}

## Files

- `index.csv` -- one row per episode
- `manifest.json` -- the picker's raw output, including all 3 human instructions
- `thumbs/<uuid>.jpg` -- 3-frame strip per episode (5% / 45% / 90%)
- `sheets/<object_class>.jpg` -- all episodes of one class stacked
- `videos/<uuid>.mp4` -- the ext1 clip

## Episodes

| uuid | object class | task | lab | frames | ok | instruction |
|---|---|---|---|---|---|---|
"""
    for t in sorted(table, key=lambda x: (x["object_class"], x["lab"])):
        ins = t["instruction_1"] or t["current_task"]
        readme += (f"| `{t['uuid']}` | {t['object_class']} | {t['task_class']} | {t['lab']} "
                   f"| {t['frames']} | {'y' if t['success'] else 'FAIL'} | {ins[:70]} |\n")

    (out / "README.md").write_text(readme)
    print(f"wrote {out}: {len(table)} episodes, {len(by_class)} class sheets")
    if missing:
        print(f"MISSING video for {len(missing)}: {missing[:5]}")


if __name__ == "__main__":
    main()