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",  # datasets encodes Image examples through PIL in the generator path8#   "pyarrow>=18",9# ]10# ///11"""Build the canonical training parquet from a bucket teacher pass -> ONE final-schema push.12 13Takes the parquet parts that falcon-perception-bucket.py wrote, joins the image bytes back in14from the source bucket (skipped when the parts already carry an `image` column, i.e. the teacher15ran with --embed-images), excludes (and ASSERTS the exclusion of) a gold slice, splits16train/validation, and writes everything in a single final-schema write -- either to a dataset17repo (one push_to_hub, never a second) or as train.parquet / validation.parquet in a bucket for18`materialize-coco.py` / direct `load_dataset` use.19 20Image bytes are fetched in chunks through a dataset generator, so RAM stays bounded to one21chunk (--chunk, default 256 images) however large the corpus is; the Arrow cache on disk holds22the rest.23 24    uv run embed-bucket-images.py \\25      --parts "hf://buckets/<ns>/<teacher-out>/part-*.parquet" \\26      --src <ns>/<source-bucket> \\27      --gold <ns>/<gold-dataset> \\28      --out <ns>/<training-dataset> --private29 30    # bucket output instead of a dataset repo:31    ... --out hf://buckets/<ns>/<training-bucket>/dataset32 33    # local everything (smoke tests, a laptop-sized corpus): --src is a directory holding the34    # keys as relative paths, --out an absolute or ./relative directory35    ... --parts "./parts/part-*.parquet" --src ./pages --gold ./gold.parquet --out ./dataset36 37Why this exists (both failure modes measured): the bucket path's output has no image column38by default, so every agent hand-writes this join; and staging an intermediate push then39re-pushing a different schema to the same repo id leaves stale repo features ->40load_dataset CastError. This script builds the final schema in memory and writes exactly once.41"""42 43import argparse44import subprocess45import tempfile46from concurrent.futures import ThreadPoolExecutor47from pathlib import Path48 49import fsspec50from datasets import ClassLabel, Dataset, DatasetDict, Image, Sequence, load_dataset51 52 53def is_local_path(s):54    # explicit prefixes only: a repo id like ns/name must never be mistaken for a directory55    # that happens to exist in the cwd56    return s.startswith(("/", "./", "../", "~"))57 58 59def main():60    p = argparse.ArgumentParser(description=__doc__.splitlines()[0])61    p.add_argument(62        "--parts",63        required=True,64        help='parquet glob, e.g. "hf://buckets/ns/out/part-*.parquet" (or a local glob)',65    )66    p.add_argument(67        "--src",68        default=None,69        help="source image bucket, e.g. ns/pages (keys = __source_key), or a local directory; "70        "not needed when the parts already carry an image column",71    )72    p.add_argument(73        "--out",74        required=True,75        help="dataset repo id, hf://buckets/... prefix, or a local directory for parquet files",76    )77    p.add_argument(78        "--gold",79        default=None,80        help="gold dataset repo id (or parquet glob / local path) to exclude, by image_id",81    )82    p.add_argument("--val-frac", type=float, default=0.1)83    p.add_argument("--seed", type=int, default=42)84    p.add_argument(85        "--limit", type=int, default=None, help="debug: cap rows AFTER gold exclusion"86    )87    p.add_argument(88        "--keep-errors",89        action="store_true",90        help="keep rows whose teacher pass errored (dropped by default: they have no usable image)",91    )92    p.add_argument(93        "--allow-gold-disjoint",94        action="store_true",95        help="permit a gold set that shares no image_id with this corpus (a genuinely different corpus)",96    )97    p.add_argument("--workers", type=int, default=16)98    p.add_argument(99        "--chunk", type=int, default=256, help="images fetched per generator chunk"100    )101    p.add_argument("--private", action="store_true")102    args = p.parse_args()103 104    try:105        ds = load_dataset("parquet", data_files=args.parts, split="train")106    except FileNotFoundError:107        raise SystemExit(108            f"no parquet files match {args.parts!r} — check the glob and bucket path"109        )110    total = len(ds)111    if total == 0:112        raise SystemExit(f"{args.parts!r} matched files but they contain 0 rows")113    if args.src and args.src.startswith("hf://buckets/"):114        args.src = args.src[len("hf://buckets/") :]115    if (116        not (args.out.startswith("hf://buckets/") or is_local_path(args.out))117        and "/" not in args.out118    ):119        raise SystemExit(120            f"--out {args.out!r}: a dataset repo id needs a namespace (ns/name)"121        )122    if "error" in ds.column_names and not args.keep_errors:123        n_err = sum(1 for e in ds["error"] if e)124        if n_err:125            ds = ds.filter(lambda r: not r["error"])126            print(f"dropped {n_err} teacher error rows (--keep-errors to keep them)")127    if not isinstance(ds.features["objects"]["category"].feature, ClassLabel):128        feats = ds.features.copy()129        feats["objects"]["category"] = Sequence(ClassLabel(names=[ds[0]["query"]]))130        ds = ds.cast(feats)131 132    # ---- gold exclusion, asserted on the stable image_id (never on path-shaped keys) ----133    if args.gold:134        if "://" in args.gold or is_local_path(args.gold):135            gold = load_dataset("parquet", data_files=args.gold, split="train")136        else:137            gold = load_dataset(args.gold, split="train")138        gold_ids = set(gold["image_id"])139        before = len(ds)140        ds = ds.filter(lambda r: r["image_id"] not in gold_ids)141        removed = before - len(ds)142        overlap = gold_ids & set(ds["image_id"])143        assert not overlap, (144            f"gold exclusion FAILED: {len(overlap)} gold ids remain, e.g. {sorted(overlap)[:3]}"145        )146        if removed == 0 and gold_ids and not args.allow_gold_disjoint:147            raise SystemExit(148                "gold exclusion matched 0 rows — the gold set and this corpus share no image_id. "149                "That usually means the ids were built from different key prefixes. If this corpus "150                "really is disjoint from the gold set, pass --allow-gold-disjoint."151            )152        print(f"gold: excluded {removed} rows; {len(gold_ids)} gold ids, overlap now 0")153 154    if args.limit:155        ds = ds.select(range(min(args.limit, len(ds))))156 157    # ---- join image bytes back in from the source (unless the parts already carry them) ----158    if "image" in ds.column_names:159        print(160            "parts already carry an image column (teacher ran with --embed-images) — no fetch"161        )162        ds = ds.cast_column("image", Image())163    else:164        if not args.src:165            raise SystemExit(166                "parts have no image column — pass --src <bucket or local dir>"167            )168        src_dir = Path(args.src).expanduser() if is_local_path(args.src) else None169 170        def fetch(key):171            if src_dir is not None:172                return (src_dir / key).read_bytes()173            with fsspec.open(f"hf://buckets/{args.src}/{key}", "rb") as f:174                return f.read()175 176        plain = ds.with_format(None)177        features = plain.features.copy()178        features["image"] = Image()179 180        def rows_with_images():181            # one chunk of bytes in RAM at a time; datasets streams the yielded rows to182            # its Arrow cache on disk, so corpus size never sets the RAM ceiling183            for start in range(0, len(plain), args.chunk):184                chunk = plain[start : start + args.chunk]  # dict of column -> list185                keys = chunk["__source_key"]186                with ThreadPoolExecutor(args.workers) as ex:187                    blobs = list(ex.map(fetch, keys))188                bad = [k for k, b in zip(keys, blobs) if not b]189                assert not bad, f"{len(bad)} images fetched empty, e.g. {bad[:3]}"190                for i, blob in enumerate(blobs):191                    row = {col: chunk[col][i] for col in chunk}192                    row["image"] = {"bytes": blob, "path": None}193                    yield row194 195        ds = Dataset.from_generator(rows_with_images, features=features)196 197    # ---- split, then ONE write ----198    if args.val_frac <= 0 or len(ds) < 2:199        out = DatasetDict({"train": ds})200        print(201            f"rows: {total} read -> {len(ds)} kept -> train only (no validation split; "202            "materialize-coco.py then needs --splits train)"203        )204    else:205        parts = ds.train_test_split(test_size=args.val_frac, seed=args.seed)206        out = DatasetDict({"train": parts["train"], "validation": parts["test"]})207        print(208            f"rows: {total} read -> {len(ds)} kept -> train {len(out['train'])} / validation {len(out['validation'])}"209        )210 211    if args.out.startswith("hf://buckets/"):212        with tempfile.TemporaryDirectory() as td:213            for split, d in out.items():214                local = Path(td) / f"{split}.parquet"215                d.to_parquet(local)216                subprocess.run(217                    ["hf", "cp", str(local), f"{args.out.rstrip('/')}/{split}.parquet"],218                    check=True,219                )220        print(f"wrote {' + '.join(f'{s}.parquet' for s in out)} -> {args.out}")221    elif is_local_path(args.out):222        out_dir = Path(args.out).expanduser()223        out_dir.mkdir(parents=True, exist_ok=True)224        for split, d in out.items():225            d.to_parquet(out_dir / f"{split}.parquet")226        print(f"wrote {' + '.join(f'{s}.parquet' for s in out)} -> {out_dir}")227    else:228        out.push_to_hub(args.out, private=args.private)229        print(f"pushed -> https://huggingface.co/datasets/{args.out}")230 231 232if __name__ == "__main__":233    main()234