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"""Materialize a COCO directory tree FROM the canonical parquet, in-job, on ephemeral disk.13 14Some trainers (RF-DETR and friends) refuse HF datasets and demand the canonical COCO 201715layout: annotations/instances_train2017.json + train2017/*.jpg. Never hand-assemble or16upload that tree -- generate it from the parquet with this script instead. A generated17tree cannot reference images that are not there, which kills the referenced-vs-uploaded18mismatch class outright (it caused three paid job failures in one measured run).19 20    # inside the training job, before the trainer starts:21    uv run materialize-coco.py --data hf://buckets/<ns>/<training-bucket>/dataset --out /tmp/coco22 23    # or from a dataset repo produced by embed-bucket-images.py:24    uv run materialize-coco.py --data <ns>/<training-dataset> --out /tmp/coco25 26    # training more than once? generate ONCE onto a bucket mount and let later jobs reuse it:27    #   hf jobs run ... -v hf://buckets/<ns>/<training-bucket>:/data ...28    uv run materialize-coco.py --data /data/dataset --out /data/coco29 30A split is reused, not regenerated, when its tree is complete AND was built from the same31labels: the annotations file carries a fingerprint of (image ids, boxes), so a corrected32dataset -- the step-6 loop, same images, new labels -- rebuilds automatically. --force rebuilds33regardless. Rows whose image cannot be decoded (teacher error rows, truncated files) are skipped34and counted, never allowed to kill the job.35 36Boxes are converted yolo-normalized -> COCO xywh pixels (pass --bbox-format coco_xywh if your37parquet already stores pixels). masks_rle, when present, is carried through as COCO RLE38segmentation (RF-DETR-class trainers accept RLE natively).39"""40 41import argparse42import hashlib43import io44import json45import shutil46from pathlib import Path47 48import numpy as np49from datasets import Image as HFImage50from datasets import load_dataset51from PIL import Image as PILImage52from pycocotools import mask as mask_utils53 54SPLIT_DIR = {"train": "train2017", "validation": "val2017"}55 56 57def to_xywh(bbox, w, h, fmt):58    if fmt == "coco_xywh":59        return [float(v) for v in bbox]60    cx, cy, bw, bh = bbox  # yolo normalised61    return [(cx - bw / 2) * w, (cy - bh / 2) * h, bw * w, bh * h]62 63 64def label_fingerprint(ds):65    """Hash of (image_id, boxes) for every row -- changes when labels change, not when bytes do."""66    labels = ds.select_columns(["image_id", "objects"]).with_format(None)67    items = sorted(68        (69            int(row["image_id"]),70            [[round(float(v), 6) for v in b] for b in row["objects"]["bbox"]],71        )72        for row in labels73    )74    return hashlib.sha1(json.dumps(items).encode()).hexdigest()75 76 77def load_split(data, split):78    if data.startswith("hf://"):79        return load_dataset(80            "parquet", data_files=f"{data.rstrip('/')}/{split}.parquet", split="train"81        )82    return load_dataset(data, split=split)83 84 85def tree_is_reusable(jpath, img_dir, fingerprint):86    if not jpath.exists():87        return False, "no tree yet"88    coco = json.loads(jpath.read_text())89    stamped = coco.get("provenance", {}).get("fingerprint")90    if stamped != fingerprint:91        return False, "labels changed since the tree was built"92    referenced = len(coco["images"])93    present = len(list(img_dir.glob("*.jpg")))94    if not referenced or referenced != present:95        return False, f"tree incomplete ({present} files vs {referenced} referenced)"96    return True, f"complete tree ({present} images), same labels"97 98 99def decode_image(raw):100    """raw is the undecoded {bytes, path} struct (or None for error rows)."""101    if not raw or not raw.get("bytes"):102        return None103    try:104        im = PILImage.open(io.BytesIO(raw["bytes"]))105        im.load()106        return im.convert("RGB")107    except Exception:  # noqa: BLE001 -- any decode failure means "skip this row"108        return None109 110 111def main():112    p = argparse.ArgumentParser(description=__doc__.splitlines()[0])113    p.add_argument(114        "--data",115        required=True,116        help="dataset repo id, hf://buckets/... prefix, or local directory holding <split>.parquet",117    )118    p.add_argument(119        "--out",120        required=True,121        help="output dir (ephemeral disk, or a bucket mount to reuse across jobs)",122    )123    p.add_argument("--bbox-format", default="yolo", choices=["yolo", "coco_xywh"])124    p.add_argument("--splits", nargs="+", default=["train", "validation"])125    p.add_argument(126        "--force",127        action="store_true",128        help="rebuild a split even if its tree is complete",129    )130    args = p.parse_args()131 132    out = Path(args.out)133    (out / "annotations").mkdir(parents=True, exist_ok=True)134 135    for split in args.splits:136        ds = load_split(args.data, split)137        assert "image" in ds.column_names, (138            "no image column — run embed-bucket-images.py first"139        )140        img_dir = out / SPLIT_DIR.get(split, split)141        jpath = out / "annotations" / f"instances_{SPLIT_DIR.get(split, split)}.json"142 143        fingerprint = label_fingerprint(ds)144        reusable, why = tree_is_reusable(jpath, img_dir, fingerprint)145        if reusable and not args.force:146            print(f"{split}: reusing {why} at {img_dir} — pass --force to rebuild")147            continue148        print(f"{split}: building ({'--force' if args.force else why})")149        # a rebuild starts from nothing: stale JPEGs from an older tree would fail the150        # files == referenced assert below after all the decode work151        shutil.rmtree(img_dir, ignore_errors=True)152        jpath.unlink(missing_ok=True)153        img_dir.mkdir()154 155        cat_feature = ds.features["objects"]["category"].feature156        names = getattr(cat_feature, "names", None) or ["object"]157 158        # undecoded bytes so a corrupt image is OUR decision to skip, not a crash inside datasets159        rows = ds.cast_column("image", HFImage(decode=False)).with_format(None)160        images, annotations, ann_id, skipped = [], [], 1, []161        for row in rows:162            iid = int(row["image_id"])163            im = None if row.get("error") else decode_image(row["image"])164            if im is None:165                skipped.append(iid)166                continue167            fname = f"{iid}.jpg"168            im.save(img_dir / fname, "JPEG", quality=95)169            # dims from the DECODED image, never metadata columns: error rows carry170            # width/height=None, and the saved JPEG is the frame everything must match171            w, h = im.size172            images.append({"id": iid, "file_name": fname, "width": w, "height": h})173            rles = json.loads(row["masks_rle"]) if row.get("masks_rle") else []174            for i, bbox in enumerate(row["objects"]["bbox"]):175                x, y, bw, bh = to_xywh(bbox, w, h, args.bbox_format)176                ann = {177                    "id": ann_id,178                    "image_id": iid,179                    "category_id": int(row["objects"]["category"][i]) + 1,180                    "bbox": [x, y, bw, bh],181                    "area": bw * bh,182                    "iscrowd": 0,183                }184                if i < len(rles):185                    rle = rles[i]186                    # masks live in the INFERENCE frame, which diverges from the image187                    # frame whenever the teacher thumbnailed -- resize before writing188                    if rle["size"] != [h, w]:189                        seg = mask_utils.decode(190                            {**rle, "counts": rle["counts"].encode()}191                        )192                        seg = np.asarray(193                            PILImage.fromarray(seg).resize((w, h), PILImage.NEAREST)194                        )195                        enc = mask_utils.encode(np.asfortranarray(seg))196                        rle = {"size": [h, w], "counts": enc["counts"].decode("ascii")}197                    ann["segmentation"] = rle198                annotations.append(ann)199                ann_id += 1200 201        coco = {202            "images": images,203            "annotations": annotations,204            "categories": [{"id": i + 1, "name": n} for i, n in enumerate(names)],205            "provenance": {206                "source": args.data,207                "fingerprint": fingerprint,208                "skipped_image_ids": skipped,209            },210        }211        jpath.write_text(json.dumps(coco))212 213        n_files = len(list(img_dir.glob("*.jpg")))214        assert n_files == len(images), (215            f"{split}: {n_files} files != {len(images)} referenced"216        )217        note = (218            f" (skipped {len(skipped)} undecodable/error rows, e.g. {skipped[:3]})"219            if skipped220            else ""221        )222        print(223            f"{split}: {len(images)} images / {len(annotations)} annotations -> {img_dir} + {jpath.name}{note}"224        )225 226 227if __name__ == "__main__":228    main()229