uv-scripts/object-detection

1#!/usr/bin/env -S uv run --script2# /// script3# requires-python = ">=3.10"4# dependencies = [5#   "datasets>=4.0",6#   "huggingface_hub>=1.27",  # hf://buckets in HfFileSystem (1.6) + prefix-collision fix (1.27)7#   "pillow",8#   "numpy",9#   "pycocotools>=2.0.11",10# ]11# ///12"""Render detection overlays from a dataset in this directory's schema -- and PROVE they rendered.13 14Draws boxes (and masks, when masks_rle is present) over the embedded images and writes PNGs.15Before reporting success it pixel-diffs every render against its source image: a page with16instances whose render is identical to the source means the overlay silently failed (alpha17bugs, empty mask lists, wrong-column reads -- all observed in real runs, twice shown to a18human as "done"). Any blank render exits nonzero and names the file.19 20    uv run render-detections.py <ns>/<teacher-or-training-dataset> --limit 10 --out previews/21    uv run render-detections.py "hf://buckets/<ns>/<bucket>/dataset/train.parquet" --out previews/22"""23 24import argparse25import io26import json27import sys28from pathlib import Path29 30import numpy as np31from datasets import Image as HFImage32from datasets import load_dataset33from PIL import Image, ImageDraw34 35COLORS = [36    (255, 210, 0),37    (80, 200, 120),38    (90, 160, 255),39    (230, 90, 80),40    (200, 120, 220),41    (255, 150, 50),42]43 44 45def main():46    p = argparse.ArgumentParser(description=__doc__.splitlines()[0])47    p.add_argument(48        "data", help="dataset repo id, or a parquet path/glob (hf:// or local)"49    )50    p.add_argument("--split", default="train")51    p.add_argument("--limit", type=int, default=10)52    p.add_argument("--out", default="previews")53    p.add_argument("--bbox-format", default="yolo", choices=["yolo", "coco_xywh"])54    p.add_argument("--no-masks", action="store_true")55    p.add_argument(56        "--min-pixels",57        type=int,58        default=1,59        help="a page with instances whose render changed fewer pixels than this is BLANK "60        "(default 1: any drawn pixel proves the overlay; a 50x50 box on a 3000px scan is real)",61    )62    args = p.parse_args()63 64    if "://" in args.data or args.data.endswith(".parquet"):65        ds = load_dataset("parquet", data_files=args.data, split="train")66    else:67        ds = load_dataset(args.data, split=args.split)68    assert "image" in ds.column_names, (69        "no image column in this dataset — nothing to render over"70    )71    ds = ds.select(range(min(args.limit, len(ds))))72 73    out = Path(args.out)74    out.mkdir(parents=True, exist_ok=True)75    blank, rendered, skipped = [], 0, []76    # undecoded bytes: a corrupt image or an error row is skipped, not a crash inside datasets77    for row in ds.cast_column("image", HFImage(decode=False)).with_format(None):78        raw = row["image"]79        try:80            src = (81                Image.open(io.BytesIO(raw["bytes"])).convert("RGB")82                if raw and raw.get("bytes")83                else None84            )85        except Exception:  # noqa: BLE001 -- any decode failure means "skip this row"86            src = None87        if src is None or row.get("error"):88            skipped.append(row["image_id"])89            continue90        im = src.copy()91        w, h = im.size92        n = len(row["objects"]["bbox"])93 94        if not args.no_masks and row.get("masks_rle"):95            from pycocotools import mask as mask_utils96 97            overlay = Image.new("RGBA", im.size, (0, 0, 0, 0))98            for i, rle in enumerate(json.loads(row["masks_rle"])):99                seg = mask_utils.decode({**rle, "counts": rle["counts"].encode()})100                if seg.shape != (h, w):101                    seg = np.asarray(Image.fromarray(seg).resize((w, h), Image.NEAREST))102                r, g, b = COLORS[i % len(COLORS)]103                tint = np.zeros((h, w, 4), np.uint8)104                tint[seg > 0] = (r, g, b, 110)105                overlay = Image.alpha_composite(overlay, Image.fromarray(tint))106            im = Image.alpha_composite(im.convert("RGBA"), overlay).convert("RGB")107 108        draw = ImageDraw.Draw(im)109        for i, bbox in enumerate(row["objects"]["bbox"]):110            if args.bbox_format == "yolo":111                cx, cy, bw, bh = bbox112                box = [113                    (cx - bw / 2) * w,114                    (cy - bh / 2) * h,115                    (cx + bw / 2) * w,116                    (cy + bh / 2) * h,117                ]118            else:119                x, y, bw, bh = bbox120                box = [x, y, x + bw, y + bh]121            draw.rectangle(box, outline=COLORS[i % len(COLORS)], width=max(3, w // 400))122 123        name = f"{row['image_id']}_{n}inst.png"124        im.save(out / name)125 126        # ---- the point of this script: prove the overlay exists ----127        changed = int(np.any(np.asarray(src) != np.asarray(im), axis=-1).sum())128        if n > 0 and changed < args.min_pixels:129            blank.append(name)130        elif n > 0:131            rendered += 1132        print(133            f"{name}: {n} instances, {changed} pixels changed ({changed / (w * h):.2%})"134        )135 136    if blank:137        sys.exit(138            f"BLANK RENDERS ({len(blank)}): {blank} — overlays did not draw; do not show these to a human."139        )140    if skipped:141        print(f"skipped {len(skipped)} undecodable/error rows, e.g. {skipped[:3]}")142    if rendered == 0:143        sys.exit(144            "No page with instances was rendered — nothing verified; increase --limit."145        )146    print(f"OK: {rendered} non-empty renders verified against source pixels -> {out}/")147 148 149if __name__ == "__main__":150    main()151