src/data/view_samples.py

"""Manifest-driven QA: a stats block, then a contact sheet.

Run this immediately after the first extractor and before writing any more
code. A sign error or an axis-order mistake in the 3D->2D projection produces
plausible-looking garbage that stays invisible until the depth numbers make no
sense a day later.

What to look for in the sheet:
  1. boxes sit on the objects, not offset or mirrored
  2. box bottoms sit at the ground contact point for cones and barriers
  3. objects lower in the frame carry smaller distances than objects near the horizon
  4. stop-sign boxes do not include the pole
  5. nuScenes and AV2 frames look like the same kind of scene

Usage:
    python -m src.data.view_samples --manifest data/unified/manifest.parquet
    python -m src.data.view_samples --per-source 8 --classes cone barrier
    python -m src.data.view_samples --source nuscenes --only-annotated --out qa/nusc.png
    python -m src.data.view_samples --split val --classes cone
"""

from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np
import pandas as pd
from PIL import Image, ImageDraw, ImageFont

from src.common import paths, schema

CLASS_COLOURS = {
    "cone": (255, 140, 0),
    "barrier": (60, 170, 255),
    "stop_sign": (80, 220, 120),
}
NEGATIVE_COLOUR = (150, 150, 150)


# ---------------------------------------------------------------------------
# Stats -- printed before anything is drawn
# ---------------------------------------------------------------------------


def print_stats(frame: pd.DataFrame) -> None:
    objects = schema.objects_only(frame)

    print("=" * 72)
    print(f"frames  : {frame['image_path'].nunique()}")
    print(f"objects : {len(objects)}")
    print(f"sensors : {sorted(frame['sensor_id'].dropna().unique())}")
    print("=" * 72)

    print("\ninstances per source x class")
    if objects.empty:
        print("  (none)")
    else:
        print(pd.crosstab(objects["source"], objects["class"]).to_string())

    print("\nbox height (px) per class")
    if objects.empty:
        print("  (none)")
    else:
        heights = (objects["y2"] - objects["y1"]).rename("height")
        print(
            heights.groupby(objects["class"])
            .describe(percentiles=[0.05, 0.25, 0.5, 0.75, 0.95])
            .to_string()
        )

    print("\ngt_distance_m per class")
    with_distance = objects[objects["gt_distance_m"].notna()]
    if with_distance.empty:
        print("  (none)")
    else:
        print(
            with_distance.groupby("class")["gt_distance_m"]
            .describe(percentiles=[0.05, 0.5, 0.95])
            .to_string()
        )

    print_warnings(frame, objects)


def print_warnings(frame: pd.DataFrame, objects: pd.DataFrame) -> None:
    warnings: list[str] = []

    if objects.empty:
        warnings.append("manifest contains no objects at all")

    if objects["gt_distance_m"].notna().sum() == 0:
        warnings.append(
            "NO gt_distance_m ANYWHERE. This manifest will train YOLO fine and "
            "leave phase 2 with nothing to score against."
        )

    non_positive = objects[objects["gt_distance_m"].notna() & (objects["gt_distance_m"] <= 0)]
    if len(non_positive):
        warnings.append(f"{len(non_positive)} objects with gt_distance_m <= 0")

    degenerate = objects[(objects["x2"] <= objects["x1"]) | (objects["y2"] <= objects["y1"])]
    if len(degenerate):
        warnings.append(f"{len(degenerate)} degenerate boxes (x2<=x1 or y2<=y1)")

    for class_name in schema.CLASSES:
        if class_name not in set(objects["class"]):
            warnings.append(f"class '{class_name}' has zero instances")

    if "split" in frame.columns and frame["split"].notna().any():
        val_cones = objects[
            (objects["split"] == "val")
            & (objects["class"] == "cone")
            & objects["gt_distance_m"].notna()
        ]
        if len(val_cones) < 200:
            warnings.append(
                f"val split has only {len(val_cones)} cone instances with distance "
                f"-- phase 2 wants at least a few hundred"
            )

    print("\nwarnings")
    if warnings:
        for warning in warnings:
            print(f"  !! {warning}")
    else:
        print("  none")
    print()


# ---------------------------------------------------------------------------
# Contact sheet
# ---------------------------------------------------------------------------


def choose_images(frame: pd.DataFrame, per_source: int, only_annotated: bool,
                  seed: int) -> list[str]:
    """Pick `per_source` image paths from each source."""
    candidates = frame
    if only_annotated:
        annotated = set(schema.objects_only(frame)["image_path"])
        candidates = frame[frame["image_path"].isin(annotated)]

    chosen: list[str] = []
    rng = np.random.default_rng(seed)
    for source in sorted(candidates["source"].unique()):
        image_paths = candidates[candidates["source"] == source]["image_path"].unique()
        if len(image_paths) == 0:
            continue
        count = min(per_source, len(image_paths))
        picked = rng.choice(image_paths, size=count, replace=False)
        chosen.extend(sorted(picked.tolist()))
    return chosen


def load_font(size: int) -> ImageFont.ImageFont:
    for candidate in (
        "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
        "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
    ):
        if Path(candidate).exists():
            return ImageFont.truetype(candidate, size)
    return ImageFont.load_default()


def render_cell(image_path: str, rows: pd.DataFrame, root: Path,
                cell_width: int) -> Image.Image | None:
    """Draw one image with its boxes, scaled to cell_width."""
    full_path = paths.resolve_image(root, image_path)
    try:
        image = Image.open(full_path).convert("RGB")
    except (FileNotFoundError, OSError) as error:
        print(f"  !! cannot open {full_path}: {error}")
        return None

    scale = cell_width / image.width
    image = image.resize((cell_width, max(1, round(image.height * scale))))
    draw = ImageDraw.Draw(image)
    font = load_font(max(12, cell_width // 45))

    objects = rows[rows["class"].notna()]
    for _, row in objects.iterrows():
        colour = CLASS_COLOURS.get(row["class"], NEGATIVE_COLOUR)
        box = [row["x1"] * scale, row["y1"] * scale, row["x2"] * scale, row["y2"] * scale]
        draw.rectangle(box, outline=colour, width=2)

        # Label exactly as the final output will: class then distance.
        label = row["class"]
        if pd.notna(row["gt_distance_m"]):
            label = f"{label} {row['gt_distance_m']:.1f}m"
        text_xy = (box[0] + 2, max(0.0, box[1] - font.size - 3))
        text_box = draw.textbbox(text_xy, label, font=font)
        draw.rectangle(text_box, fill=colour)
        draw.text(text_xy, label, fill=(0, 0, 0), font=font)

    caption = f"{image_path}  [{len(objects)} obj]"
    caption_box = draw.textbbox((4, 4), caption, font=font)
    draw.rectangle(caption_box, fill=(0, 0, 0))
    draw.text((4, 4), caption, fill=(255, 255, 255), font=font)
    return image


def tile(cells: list[Image.Image], columns: int, padding: int = 6) -> Image.Image:
    cell_width = max(cell.width for cell in cells)
    cell_height = max(cell.height for cell in cells)
    rows = (len(cells) + columns - 1) // columns

    sheet = Image.new(
        "RGB",
        (columns * cell_width + (columns + 1) * padding,
         rows * cell_height + (rows + 1) * padding),
        (25, 25, 25),
    )
    for index, cell in enumerate(cells):
        column, row = index % columns, index // columns
        sheet.paste(
            cell,
            (padding + column * (cell_width + padding),
             padding + row * (cell_height + padding)),
        )
    return sheet


# ---------------------------------------------------------------------------


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--manifest", type=Path, default=None,
                        help="defaults to <unified-root>/manifest.parquet")
    parser.add_argument("--unified-root", type=Path, default=None)
    parser.add_argument("--per-source", type=int, default=4)
    parser.add_argument("--classes", nargs="+", default=None,
                        help="keep only these classes, and only frames containing them")
    parser.add_argument("--source", default=None, help="restrict to one source")
    parser.add_argument("--split", default=None, help="restrict to train or val")
    parser.add_argument("--only-annotated", action="store_true",
                        help="never sample a frame with no objects")
    parser.add_argument("--out", type=Path, default=Path("qa/samples.png"))
    parser.add_argument("--cell-width", type=int, default=640)
    parser.add_argument("--cols", type=int, default=4)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--stats-only", action="store_true")
    args = parser.parse_args()

    root = paths.unified_root(args.unified_root)
    manifest_path = args.manifest or (root / "manifest.parquet")
    frame = schema.read_manifest(manifest_path)
    print(f"loaded {manifest_path}  ({len(frame)} rows)\n")

    if args.source:
        frame = frame[frame["source"] == args.source]
    if args.split:
        frame = frame[frame["split"] == args.split]
    if args.classes:
        # Keep the requested classes, and only frames that contain one.
        keep = frame["class"].isin(args.classes)
        frame = frame[frame["image_path"].isin(frame[keep]["image_path"])]
        frame = frame[keep | frame["class"].isna()]

    if frame.empty:
        print("!! nothing matches those filters")
        return

    print_stats(frame)
    if args.stats_only:
        return

    image_paths = choose_images(frame, args.per_source, args.only_annotated, args.seed)
    if not image_paths:
        print("!! no images to draw")
        return

    by_image = {path: group for path, group in frame.groupby("image_path")}
    cells = []
    for image_path in image_paths:
        cell = render_cell(image_path, by_image[image_path], root, args.cell_width)
        if cell is not None:
            cells.append(cell)

    if not cells:
        print("!! every image failed to load -- check the images/<source> symlink")
        return

    sheet = tile(cells, args.cols)
    args.out.parent.mkdir(parents=True, exist_ok=True)
    sheet.save(args.out)
    print(f"wrote {args.out}  ({len(cells)} frames, {sheet.width}x{sheet.height})")


if __name__ == "__main__":
    main()