uv-scripts/object-detection

1#!/usr/bin/env -S uv run --script2# /// script3# requires-python = ">=3.10"4# dependencies = [5#   "falcon-perception>=1.0.0",6#   # tarball not git+: some GPU images have no `git` for uv to shell out to7#   "bucketbag @ https://github.com/davanstrien/bucketbag/archive/refs/tags/v0.3.1.tar.gz",8#   "pyarrow>=18",9#   "pycocotools>=2.0.11",10# ]11# ///12"""Falcon-Perception over a whole HF bucket, resumable.13 14    hf jobs uv run --flavor a10g-large --secrets HF_TOKEN \15        https://huggingface.co/datasets/uv-scripts/object-detection/raw/main/falcon-perception-bucket.py \16        --src biglam/bl-images --prefix full/embellishments \17        --out davanstrien/bl-masks --query illustration18 19Input  : bucketbag batched_files — bounded scratch, files deleted as the loop advances20Engine : PagedInferenceEngine (CUDA, continuous batching)21Output : one parquet per batch -> out bucket; resume via completed_keys(__source_key)22 23Kill it at any point and re-run the same command. Done keys are skipped.24 25Output is parquet parts in a BUCKET, not a dataset repo — that is what makes the26run resumable (`completed_keys` reads the done-set back from `__source_key`).27To hand the result to the rest of this directory, publish it once at the end:28 29    from datasets import ClassLabel, Image, Sequence, load_dataset30    ds = load_dataset("parquet", data_files="hf://buckets/<namespace>/<bucket>/part-*.parquet",31                      split="train")32    feats = ds.features.copy()  # parquet stores category as bare ints; name the class33    feats["objects"]["category"] = Sequence(ClassLabel(names=[ds[0]["query"]]))34    if "image" in feats:  # --embed-images parts: make the bytes a decodable Image column35        feats["image"] = Image()36    ds.cast(feats).push_to_hub("<namespace>/<dataset>")  # a dataset repo, distinct from the bucket37 38    uv run validate-hf-dataset.py <namespace>/<dataset> --bbox-format yolo39 40By default the parts carry `width`/`height` but no `image` column: the images stay in41the source bucket, the parts stay small and resumable, and `embed-bucket-images.py`42joins the bytes back in. Pass --embed-images to write the source bytes into each part43instead (an `image` column datasets decodes directly) -- storage is cheap and it saves44the join's re-fetch of every image; the cost is a copy of the corpus in the output bucket.45 46GOTCHAS (all measured, none in the model card):47  * --query is a CLASS NAME. "illustration" works; "the illustration, excluding48    captions" returns nothing.49  * torch.compile breaks on per-image dynamic shapes -> compile is OFF here.50  * engine_config_for_gpu() sizes from the GPU and ignores host RAM; the 15 GB51    flavors (t4-small, a10g-small) get OOMKilled (exit 137) before processing52    anything -- pick >15 GB `ram` from `hf jobs hardware --json`.53    cudagraph is off by default here for the same reason.54  * xy in the output is the NORMALISED CENTRE, not a corner.55"""56 57import argparse58import hashlib59import io60import json61import time62 63import fsspec64import pyarrow as pa65import pyarrow.parquet as pq66from bucketbag import batched_files, boost, completed_keys, iter_keys, put_files67from pycocotools import mask as mask_utils68 69 70def stable_id(key):71    """Deterministic int64 image_id from the source key (COCO consumers need an int;72    a hash, not a running index, keeps ids identical across resumed runs)."""73    return int.from_bytes(hashlib.blake2b(str(key).encode(), digest_size=8).digest(), "big") >> 174 75# Same YOLO column layout as falcon-perception.py, so both outputs validate with76# `validate-hf-dataset.py --bbox-format yolo` and can be concatenated.77# `__source_key` is bucketbag's resume column — the name is load-bearing.78SCHEMA = pa.schema([79    ("__source_key", pa.string()),80    ("image_id", pa.int64()),  # int, not str: COCO-style trainers tensorise it81    ("width", pa.int32()),82    ("height", pa.int32()),83    ("objects", pa.struct([84        ("bbox", pa.list_(pa.list_(pa.float32()))),   # yolo: cx, cy, w, h normalised85        ("category", pa.list_(pa.int64())),           # single class per run; the class NAME86                                                      # is the `query` column — cast to87                                                      # ClassLabel at publish (see docstring)88        ("area", pa.list_(pa.float32())),89        ("rectangularity", pa.list_(pa.float32())),   # triage proxy — no confidence score exists90    ])),91    ("n_instances", pa.int32()),92    ("masks_rle", pa.string()),93    ("query", pa.string()),94    ("gen_seconds", pa.float32()),95    ("error", pa.string()),96])97 98 99def pair_bboxes(raw):100    boxes, cur = [], {}101    for e in raw:102        if not isinstance(e, dict):103            continue104        cur.update(e)105        if all(k in cur for k in ("x", "y", "h", "w")):106            boxes.append(dict(cur)); cur = {}107    return boxes108 109 110def serialise(rows, fmt, schema=SCHEMA):111    if fmt == "jsonl":112        return "\n".join(json.dumps(r) for r in rows) + "\n"113    buf = io.BytesIO()114    pq.write_table(pa.Table.from_pylist(rows, schema=schema), buf, compression="zstd")115    return buf.getvalue()116 117 118def main():119    p = argparse.ArgumentParser()120    p.add_argument("--src", required=True, help="source bucket, e.g. biglam/bl-images")121    p.add_argument("--prefix", default=None, help="bucket prefix, e.g. full/embellishments")122    p.add_argument("--out", required=True, help="output bucket")123    p.add_argument("--query", default="illustration", help="a CLASS NAME, not an instruction")124    p.add_argument("--task", default="segmentation", choices=["segmentation", "detection"])125    p.add_argument("--limit", type=int, default=None)126    p.add_argument("--max-dim", type=int, default=1024)127    p.add_argument("--max-new-tokens", type=int, default=200)128    p.add_argument("--batch-n", type=int, default=32, help="files per bucketbag batch")129    p.add_argument("--max-bytes", type=int, default=None,130                   help="scratch bytes per batch (RAM tmpfs). Default 2 GiB; 256 MiB with --embed-images, "131                        "whose bytes are also held ~4x in host RAM while a part is serialised")132    p.add_argument("--cudagraph", action="store_true", help="opt IN; off by default (host OOM)")133    p.add_argument("--format", default="parquet", choices=["parquet", "jsonl"])134    p.add_argument("--no-resume", action="store_true")135    p.add_argument("--embed-images", action="store_true",136                   help="also write the source image bytes into each part (see docstring)")137    args = p.parse_args()138    if args.embed_images and args.format == "jsonl":139        raise SystemExit("--embed-images writes raw image bytes, which jsonl cannot carry; use --format parquet.")140    if args.max_bytes is None:141        args.max_bytes = 256 * 2**20 if args.embed_images else 2 * 2**30142    schema = SCHEMA143    if args.embed_images:144        schema = SCHEMA.append(pa.field("image", pa.struct([("bytes", pa.binary()), ("path", pa.string())])))145 146    if args.format == "jsonl" and not args.no_resume:147        # completed_keys only reads the done-set back from .parquet parts, so jsonl148        # output silently reprocesses EVERYTHING on every re-run.149        raise SystemExit("--format jsonl is not resumable; pass --no-resume to run it anyway.")150 151    try:152        import torch153 154        has_cuda = torch.cuda.is_available()155    except ImportError:  # falcon-perception pins torch off-darwin, so it may be absent156        has_cuda = False157    if not has_cuda:158        raise SystemExit(159            "This script needs a CUDA GPU (PagedInferenceEngine). "160            "For MLX/CPU-capable runs use falcon-perception.py instead."161        )162 163    boost()  # raise xet small-file concurrency — the whole point on many small objects164 165    from huggingface_hub import HfApi166 167    # first run: the out bucket may not exist yet — completed_keys 404s on a168    # missing bucket, killing the job before anything happens169    HfApi().create_bucket(args.out, private=True, exist_ok=True)170 171    done = set() if args.no_resume else completed_keys(args.out)172    print(f"{len(done)} keys already done", flush=True)173    if done and args.format == "parquet":174        # a resume must not mix part schemas: half the parts with an image column and half175        # without loads as nulls downstream, and the null rows crash the trainer's tree build176        first_part = next((f for f in iter_keys(args.out, prefix="part-", objects=True)177                           if f.path.endswith(".parquet")), None)178        if first_part is not None:179            with fsspec.open(f"hf://buckets/{args.out}/{first_part.path}", "rb") as fh:180                existing = pq.read_schema(fh)181            if ("image" in existing.names) != args.embed_images:182                raise SystemExit(183                    f"existing parts in {args.out} were written "184                    f"{'with' if 'image' in existing.names else 'without'} --embed-images; "185                    "resume with the same flag, or write to a fresh --out bucket."186                )187 188    # objects=True yields BucketFile (with .size), so max_bytes is honoured.189    # Needs bucketbag >= 0.3.0: before that, string keys made batched_files drop190    # max_bytes silently and run unbounded against RAM-tmpfs scratch.191    keys = [192        f for f in iter_keys(args.src, prefix=args.prefix, objects=True)193        if f.path.lower().endswith((".jpg", ".jpeg", ".png")) and f.path not in done194    ]195    if args.limit:196        keys = keys[: args.limit]197    print(f"{len(keys)} keys to process", flush=True)198    if not keys:199        raise SystemExit(200            f"0 keys matched under {args.src}/{args.prefix or ''} — this script reads only "201            ".jpg/.jpeg/.png (convert JPEG 2000 / TIFF first), and already-done keys are skipped "202            "(pass --no-resume to redo)."203        )204 205    from falcon_perception import PERCEPTION_MODEL_ID, build_prompt_for_task, load_and_prepare_model, setup_torch_config206    from falcon_perception.data import ImageProcessor207    from falcon_perception.paged_inference import (208        PagedInferenceEngine, SamplingParams, Sequence, engine_config_for_gpu,209    )210 211    setup_torch_config()212    t = time.perf_counter()213    model, tokenizer, _ = load_and_prepare_model(214        hf_model_id=PERCEPTION_MODEL_ID, dtype="bfloat16", compile=False,  # compile breaks on dynamic shapes215    )216    print(f"model loaded in {time.perf_counter() - t:.1f}s", flush=True)217 218    cfg = engine_config_for_gpu(max_image_size=args.max_dim, dtype=model.dtype)219    print(f"paged config: {cfg}", flush=True)220    engine = PagedInferenceEngine(221        model, tokenizer, ImageProcessor(patch_size=16, merge_size=1),222        max_seq_length=8192, capture_cudagraph=args.cudagraph, **cfg,223    )224    sp = SamplingParams(225        args.max_new_tokens,226        stop_token_ids=[tokenizer.eos_token_id, tokenizer.end_of_query_token_id],227        coord_dedup_threshold=0.01,228    )229    prompt = build_prompt_for_task(args.query, args.task)230 231    n, gen_total, t_all, batch_i = 0, 0.0, time.perf_counter(), 0232    for batch in batched_files(args.src, keys=keys, n=args.batch_n, max_bytes=args.max_bytes):233        # NOTE: never hold a LoadedItem past its batch — convert eagerly.234        pairs = []235        for it in batch:236            try:237                img = it.image.convert("RGB")  # convert() forces the load off disk238                orig_size = img.size  # SOURCE dims -- the images downstream tools decode239                if max(img.size) > args.max_dim * 2:240                    img.thumbnail((args.max_dim * 2, args.max_dim * 2))241                raw = it.bytes if args.embed_images else None  # read before the batch is deleted242                pairs.append((str(it.key), img, orig_size, raw))243            except Exception as e:244                pairs.append((str(it.key), e, None, None))245 246        good = [(k, im, sz, raw) for k, im, sz, raw in pairs if not isinstance(im, Exception)]247        seqs = [248            Sequence(text=prompt, image=im, min_image_size=256,249                     max_image_size=args.max_dim, request_idx=i, task=args.task)250            for i, (_, im, _, _) in enumerate(good)251        ]252        t0 = time.perf_counter()253        if seqs:254            engine.generate(seqs, sampling_params=sp)255        dt = time.perf_counter() - t0256        gen_total += dt257 258        rows = []259        for (key, im, orig_size, raw), seq in zip(good, seqs):260            aux = seq.output_aux261            boxes = pair_bboxes(aux.bboxes_raw)262            masks = list(aux.masks_rle)263            for m in masks:264                if isinstance(m.get("counts"), bytes):265                    m["counts"] = m["counts"].decode()266            # width/height are the SOURCE image's dims: boxes are normalised (frame-free),267            # and downstream pixel conversions run against the untouched bucket images.268            W, H = orig_size269            bbox, area, rect = [], [], []270            for i, b in enumerate(boxes):271                bbox.append([b["x"], b["y"], b["w"], b["h"]])  # yolo: cx, cy, w, h normalised272                a = b["w"] * b["h"]273                area.append(a)274                r = 0.0275                if i < len(masks):  # rectangularity — the only triage signal; no score exists276                    try:277                        m = masks[i]278                        if isinstance(m.get("counts"), str):279                            m = {**m, "counts": m["counts"].encode()}280                        # box area measured in the MASK's own frame (rle size) — mixing281                        # frames skews r282                        mh, mw = (m.get("size") or [H, W])[:2]283                        r = min(float(mask_utils.area(m)) / max(a * mw * mh, 1.0), 1.0)284                    except Exception:285                        r = 0.0286                rect.append(r)287            row = {288                "__source_key": key, "image_id": stable_id(key), "width": W, "height": H,289                "objects": {"bbox": bbox, "category": [0] * len(bbox),290                            "area": area, "rectangularity": rect},291                "n_instances": len(bbox), "masks_rle": json.dumps(masks),292                "query": args.query, "gen_seconds": dt / max(len(seqs), 1), "error": None,293            }294            if args.embed_images:295                row["image"] = {"bytes": raw, "path": None}296            rows.append(row)297        for key, err, _, _ in [t for t in pairs if isinstance(t[1], Exception)]:298            # a durable error row, never a gap — and it counts as done so it is299            # not retried forever on every re-run. objects is EMPTY, not null:300            # a null struct crashes validate-hf-dataset.py after publish.301            rows.append({k: None for k in SCHEMA.names} | {302                "__source_key": key, "image_id": stable_id(key), "query": args.query,303                "objects": {"bbox": [], "category": [], "area": [], "rectangularity": []},304                "n_instances": 0, "masks_rle": "[]",305                "error": f"{type(err).__name__}: {err}",306            })307 308        # part name derives from batch CONTENT, not a run-local counter: a resumed run's309        # counter restarts at 0 and put_files overwrites, silently destroying the first310        # run's parts. A content-derived name is stable per batch and collision-free311        # across resumes (a re-run of the same batch overwrites its own part, idempotent).312        if not rows:  # bucketbag drops files that vanished between listing and download313            continue314        ext = "jsonl" if args.format == "jsonl" else "parquet"315        part = hashlib.blake2b(rows[0]["__source_key"].encode(), digest_size=6).hexdigest()316        put_files([(f"part-{part}.{ext}", serialise(rows, args.format, schema))], args.out)317        n += len(rows); batch_i += 1318        rate = n / (time.perf_counter() - t_all)319        print(f"batch {batch_i}: {len(rows)} rows ({dt / max(len(seqs), 1):.2f}s/img)  "320              f"total {n}  {rate:.2f} img/s", flush=True)321 322    wall = time.perf_counter() - t_all323    print(f"\n{n} images in {wall:.1f}s ({gen_total:.1f}s generation) | {n / wall:.2f} img/s", flush=True)324    if n:325        print(f"extrapolation: 100k images ≈ {wall / n * 100_000 / 3600:.1f} GPU-hours end-to-end", flush=True)326 327 328main()329