tools/qa_sheets.py

#!/usr/bin/env python3
"""Render one contact sheet per task so the conditions can actually be looked at.

``verify_v1.py`` proves a task dir is *structurally* consistent — matching id sets, matching frame
counts, intact hashes. It cannot see that a re-warp filled the border with the wrong colour, that
an inpainting mask no longer covers its own hole, that a depth map came out flat, or that a style
reference got paired with the wrong sample. Those are all visible in a second and invisible to
every check that does not draw pixels.

Layout: one row per sample, one column per condition, three frames (first / middle / last) for
each video condition and the single image for each still condition, labelled.

Writes ``datasets/v1/qa/<root>__<task>.jpg``.

Usage
=====
    python qa_sheets.py                          # every task, up to --max-samples rows each
    python qa_sheets.py --task camera_inpainting --max-samples 20
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import cv2
import numpy as np

ROOT = Path(__file__).resolve().parents[1]
V1 = ROOT / "datasets/v1"
QA = V1 / "qa"

CELL_W, CELL_H = 210, 122
ORDER = ["camera", "pose", "depth", "inpainting", "inpainting/bbox", "subject", "style"]


def label(tile: np.ndarray, text: str) -> np.ndarray:
    cv2.rectangle(tile, (0, 0), (tile.shape[1], 15), (0, 0, 0), -1)
    cv2.putText(tile, text[:34], (3, 11), cv2.FONT_HERSHEY_SIMPLEX, 0.33, (255, 255, 255), 1)
    return tile


def blank(text: str = "") -> np.ndarray:
    tile = np.full((CELL_H, CELL_W, 3), 28, np.uint8)
    return label(tile, text) if text else tile


def video_strip(path: Path, name: str) -> np.ndarray:
    cap = cv2.VideoCapture(str(path))
    frames = []
    while True:
        ok, frame = cap.read()
        if not ok:
            break
        frames.append(frame)
    cap.release()
    if not frames:
        return np.hstack([blank(f"{name} EMPTY")] * 3)
    picks = [frames[0], frames[len(frames) // 2], frames[-1]]
    tiles = []
    for i, frame in enumerate(picks):
        tile = cv2.resize(frame, (CELL_W, CELL_H))
        tiles.append(label(tile, f"{name} [{['first','mid','last'][i]}]" if i == 0 else
                           f"{['','mid','last'][i]}"))
    return np.hstack(tiles)


def image_cell(path: Path, name: str) -> np.ndarray:
    image = cv2.imread(str(path))
    if image is None:
        return np.hstack([blank(f"{name} MISSING")] * 3)
    scale = min(CELL_W * 3 / image.shape[1], CELL_H / image.shape[0])
    resized = cv2.resize(image, (max(1, int(image.shape[1] * scale)), max(1, int(image.shape[0] * scale))))
    canvas = np.full((CELL_H, CELL_W * 3, 3), 28, np.uint8)
    y, x = (CELL_H - resized.shape[0]) // 2, (CELL_W * 3 - resized.shape[1]) // 2
    canvas[y:y + resized.shape[0], x:x + resized.shape[1]] = resized
    return label(canvas, name)


def sheet_for(task_dir: Path, max_samples: int) -> np.ndarray | None:
    pairing = json.loads((task_dir / "pairing.json").read_text())
    conditions = [c for c in ORDER if c in pairing["conditions"]]
    if "inpainting" in conditions:
        conditions.insert(conditions.index("inpainting") + 1, "inpainting/bbox")
    samples = sorted(pairing["samples"])[:max_samples]
    if not samples:
        return None

    rows = []
    for sid in samples:
        cells = []
        for cond in conditions:
            if cond in {"subject", "style"}:
                found = next((p for p in (task_dir / cond).glob(f"{sid}.*")), None)
                cells.append(image_cell(found, f"{sid} {cond}") if found else blank(f"{cond} MISSING"))
            else:
                path = task_dir / cond / f"{sid}.mp4"
                cells.append(video_strip(path, f"{sid} {cond}") if path.exists()
                             else np.hstack([blank(f"{cond} MISSING")] * 3))
        rows.append(np.hstack(cells))
    width = max(r.shape[1] for r in rows)
    rows = [np.pad(r, ((0, 2), (0, width - r.shape[1]), (0, 0))) for r in rows]
    return np.vstack(rows)


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--task", action="append", default=None)
    ap.add_argument("--max-samples", type=int, default=6)
    args = ap.parse_args()

    QA.mkdir(parents=True, exist_ok=True)
    task_dirs = sorted(p.parent for p in V1.glob("pairwise/*/pairing.json"))
    task_dirs += sorted(p.parent for p in V1.glob("more_conditions/*/*/pairing.json"))
    if args.task:
        wanted = set(args.task)
        task_dirs = [d for d in task_dirs if d.name in wanted]

    for task_dir in task_dirs:
        sheet = sheet_for(task_dir, args.max_samples)
        if sheet is None:
            continue
        out = QA / (str(task_dir.relative_to(V1)).replace("/", "__") + ".jpg")
        cv2.imwrite(str(out), sheet, [cv2.IMWRITE_JPEG_QUALITY, 82])
        print(f"  {out.relative_to(V1)}  ({sheet.shape[1]}x{sheet.shape[0]})")
    print(f"\nWrote {len(task_dirs)} sheet(s) to {QA}")


if __name__ == "__main__":
    main()