uv-scripts/object-detection

1# /// script2# requires-python = ">=3.11"3# dependencies = [4#     "datasets>=3.1.0",5#     "huggingface-hub",6#     "tqdm",7#     "Pillow", 8# ]9# ///10 11"""12Generate rich statistics for object detection datasets on Hugging Face Hub.13 14Mirrors panlabel's stats command. Computes:15 16- Summary counts (images, annotations, categories)17- Label distribution histogram (top-N)18- Bounding box statistics (area, aspect ratio, out-of-bounds)19- Annotation density per image20- Per-category bbox statistics21- Category co-occurrence pairs22- Image resolution distribution23 24Supports COCO-style (xywh), XYXY/VOC, YOLO (normalized center xywh),25TFOD (normalized xyxy), and Label Studio (percentage xywh) bbox formats.26Supports streaming for large datasets. Outputs text or JSON.27 28Examples:29  uv run stats-hf-dataset.py merve/test-coco-dataset30  uv run stats-hf-dataset.py merve/test-coco-dataset --top 20 --report json31  uv run stats-hf-dataset.py merve/test-coco-dataset --bbox-format tfod32  uv run stats-hf-dataset.py merve/test-coco-dataset --streaming --max-samples 500033"""34 35import argparse36import json37import logging38import math39import os40import sys41import time42from collections import Counter, defaultdict43from datetime import datetime44from typing import Any45 46from datasets import load_dataset47from huggingface_hub import DatasetCard, login48from tqdm.auto import tqdm49 50logging.basicConfig(level=logging.INFO)51logger = logging.getLogger(__name__)52 53BBOX_FORMATS = ["coco_xywh", "xyxy", "voc", "yolo", "tfod", "label_studio"]54 55 56def to_xyxy(bbox: list[float], fmt: str, img_w: float = 1.0, img_h: float = 1.0) -> tuple[float, float, float, float]:57    """Convert any bbox format to (xmin, ymin, xmax, ymax) in pixel space."""58    if fmt == "coco_xywh":59        x, y, w, h = bbox60        return (x, y, x + w, y + h)61    elif fmt in ("xyxy", "voc"):62        return tuple(bbox[:4])63    elif fmt == "yolo":64        cx, cy, w, h = bbox65        return (cx - w / 2) * img_w, (cy - h / 2) * img_h, (cx + w / 2) * img_w, (cy + h / 2) * img_h66    elif fmt == "tfod":67        xmin_n, ymin_n, xmax_n, ymax_n = bbox68        return (xmin_n * img_w, ymin_n * img_h, xmax_n * img_w, ymax_n * img_h)69    elif fmt == "label_studio":70        x_pct, y_pct, w_pct, h_pct = bbox71        return (72            x_pct / 100.0 * img_w,73            y_pct / 100.0 * img_h,74            (x_pct + w_pct) / 100.0 * img_w,75            (y_pct + h_pct) / 100.0 * img_h,76        )77    else:78        raise ValueError(f"Unknown bbox format: {fmt}")79 80 81def percentile(sorted_vals: list[float], p: float) -> float:82    """Compute percentile from sorted values."""83    if not sorted_vals:84        return 0.085    k = (len(sorted_vals) - 1) * p / 100.086    f = int(k)87    c = f + 188    if c >= len(sorted_vals):89        return sorted_vals[-1]90    return sorted_vals[f] + (k - f) * (sorted_vals[c] - sorted_vals[f])91 92 93def main(94    input_dataset: str,95    bbox_column: str = "bbox",96    category_column: str = "category",97    bbox_format: str = "coco_xywh",98    width_column: str | None = "width",99    height_column: str | None = "height",100    split: str = "train",101    max_samples: int | None = None,102    streaming: bool = False,103    top: int = 10,104    report_format: str = "text",105    tolerance: float = 0.5,106    hf_token: str | None = None,107    output_dataset: str | None = None,108    private: bool = False,109):110    """Compute statistics for an object detection dataset."""111 112    start_time = datetime.now()113 114    HF_TOKEN = hf_token or os.environ.get("HF_TOKEN")115    if HF_TOKEN:116        login(token=HF_TOKEN)117 118    logger.info(f"Loading dataset: {input_dataset} (split={split}, streaming={streaming})")119    dataset = load_dataset(input_dataset, split=split, streaming=streaming)120 121    # Accumulators122    total_images = 0123    total_annotations = 0124    category_counts = Counter()125    annotations_per_image = []126    areas = []127    aspect_ratios = []128    widths = []129    heights = []130    out_of_bounds_count = 0131    zero_area_count = 0132    per_category_areas = defaultdict(list)133    co_occurrence_pairs = Counter()134    images_without_annotations = 0135 136    iterable = dataset137    if max_samples:138        if streaming:139            iterable = dataset.take(max_samples)140        else:141            iterable = dataset.select(range(min(max_samples, len(dataset))))142 143    for idx, example in enumerate(tqdm(iterable, desc="Computing stats", total=max_samples)):144        total_images += 1145 146        objects = example.get("objects", example)147        bboxes = objects.get(bbox_column, []) or []148        categories = objects.get(category_column, []) or []149 150        # Image dimensions151        img_w = None152        img_h = None153        if width_column:154            img_w = example.get(width_column) or (objects.get(width_column) if isinstance(objects, dict) else None)155        if height_column:156            img_h = example.get(height_column) or (objects.get(height_column) if isinstance(objects, dict) else None)157 158        if img_w is not None and img_h is not None:159            widths.append(img_w)160            heights.append(img_h)161 162        num_anns = len(bboxes)163        annotations_per_image.append(num_anns)164        total_annotations += num_anns165 166        if num_anns == 0:167            images_without_annotations += 1168            continue169 170        # Track categories and co-occurrences171        image_cats = set()172        for ann_idx, bbox in enumerate(bboxes):173            cat = categories[ann_idx] if ann_idx < len(categories) else None174            cat_str = str(cat) if cat is not None else "<unknown>"175            category_counts[cat_str] += 1176            image_cats.add(cat_str)177 178            if bbox is None or len(bbox) < 4:179                continue180            if not all(math.isfinite(v) for v in bbox[:4]):181                continue182 183            w_for_conv = img_w if img_w else 1.0184            h_for_conv = img_h if img_h else 1.0185            xmin, ymin, xmax, ymax = to_xyxy(bbox[:4], bbox_format, w_for_conv, h_for_conv)186 187            bw = xmax - xmin188            bh = ymax - ymin189            area = bw * bh190 191            if area <= 0:192                zero_area_count += 1193            else:194                areas.append(area)195                per_category_areas[cat_str].append(area)196 197            if bh > 0:198                aspect_ratios.append(bw / bh)199 200            # Out of bounds check201            if img_w is not None and img_h is not None:202                if xmin < -tolerance or ymin < -tolerance or xmax > img_w + tolerance or ymax > img_h + tolerance:203                    out_of_bounds_count += 1204 205        # Co-occurrence pairs206        sorted_cats = sorted(image_cats)207        for i in range(len(sorted_cats)):208            for j in range(i + 1, len(sorted_cats)):209                co_occurrence_pairs[(sorted_cats[i], sorted_cats[j])] += 1210 211    processing_time = datetime.now() - start_time212 213    # Compute distribution stats214    areas.sort()215    aspect_ratios.sort()216    annotations_per_image.sort()217 218    def dist_stats(vals: list[float]) -> dict:219        if not vals:220            return {"count": 0, "min": 0, "max": 0, "mean": 0, "median": 0, "p25": 0, "p75": 0}221        return {222            "count": len(vals),223            "min": round(vals[0], 2),224            "max": round(vals[-1], 2),225            "mean": round(sum(vals) / len(vals), 2),226            "median": round(percentile(vals, 50), 2),227            "p25": round(percentile(vals, 25), 2),228            "p75": round(percentile(vals, 75), 2),229        }230 231    # Top-N categories232    top_categories = category_counts.most_common(top)233 234    # Top co-occurrence pairs235    top_cooccurrences = co_occurrence_pairs.most_common(top)236 237    # Per-category bbox area stats238    per_cat_stats = {}239    for cat, cat_areas in sorted(per_category_areas.items(), key=lambda x: -len(x[1])):240        cat_areas.sort()241        per_cat_stats[cat] = dist_stats(cat_areas)242 243    report = {244        "dataset": input_dataset,245        "split": split,246        "summary": {247            "total_images": total_images,248            "total_annotations": total_annotations,249            "unique_categories": len(category_counts),250            "images_without_annotations": images_without_annotations,251            "out_of_bounds_bboxes": out_of_bounds_count,252            "zero_area_bboxes": zero_area_count,253        },254        "label_distribution": {cat: count for cat, count in top_categories},255        "annotation_density": dist_stats([float(x) for x in annotations_per_image]),256        "bbox_area": dist_stats(areas),257        "bbox_aspect_ratio": dist_stats(aspect_ratios),258        "image_resolution": {259            "width": dist_stats([float(w) for w in sorted(widths)]) if widths else {},260            "height": dist_stats([float(h) for h in sorted(heights)]) if heights else {},261        },262        "per_category_area": {cat: per_cat_stats[cat] for cat in list(per_cat_stats)[:top]},263        "co_occurrence_pairs": [264            {"pair": list(pair), "count": count} for pair, count in top_cooccurrences265        ],266        "processing_time_seconds": processing_time.total_seconds(),267        "timestamp": datetime.now().isoformat(),268    }269 270    if report_format == "json":271        print(json.dumps(report, indent=2))272    else:273        print("\n" + "=" * 60)274        print(f"Dataset Statistics: {input_dataset}")275        print("=" * 60)276 277        s = report["summary"]278        print(f"\n  Images:           {s['total_images']:,}")279        print(f"  Annotations:      {s['total_annotations']:,}")280        print(f"  Categories:       {s['unique_categories']:,}")281        print(f"  Empty images:     {s['images_without_annotations']:,}")282        print(f"  Out-of-bounds:    {s['out_of_bounds_bboxes']:,}")283        print(f"  Zero-area bboxes: {s['zero_area_bboxes']:,}")284 285        if total_images > 0:286            print(f"\n  Annotations/image: {total_annotations / total_images:.1f} avg")287 288        d = report["annotation_density"]289        if d["count"]:290            print(f"    min={d['min']}, median={d['median']}, max={d['max']}")291 292        print(f"\n  Label Distribution (top {top}):")293        for cat, count in top_categories:294            pct = 100.0 * count / total_annotations if total_annotations else 0295            bar = "#" * int(pct / 2)296            print(f"    {cat:30s} {count:>8,} ({pct:5.1f}%) {bar}")297 298        a = report["bbox_area"]299        if a["count"]:300            print(f"\n  Bbox Area:")301            print(f"    min={a['min']}, median={a['median']}, mean={a['mean']}, max={a['max']}")302 303        ar = report["bbox_aspect_ratio"]304        if ar["count"]:305            print(f"\n  Bbox Aspect Ratio (w/h):")306            print(f"    min={ar['min']}, median={ar['median']}, mean={ar['mean']}, max={ar['max']}")307 308        if top_cooccurrences:309            print(f"\n  Category Co-occurrence (top {top}):")310            for pair, count in top_cooccurrences:311                print(f"    {pair[0]} + {pair[1]}: {count:,}")312 313        print(f"\n  Processing time: {processing_time.total_seconds():.1f}s")314        print("=" * 60)315 316    # Optionally push stats report as a dataset317    if output_dataset:318        from datasets import Dataset as HFDataset319 320        report_ds = HFDataset.from_dict({321            "report_json": [json.dumps(report)],322            "dataset": [input_dataset],323            "total_images": [total_images],324            "total_annotations": [total_annotations],325            "unique_categories": [len(category_counts)],326            "timestamp": [datetime.now().isoformat()],327        })328 329        logger.info(f"Pushing stats report to {output_dataset}")330        max_retries = 3331        for attempt in range(1, max_retries + 1):332            try:333                if attempt > 1:334                    os.environ["HF_HUB_DISABLE_XET"] = "1"335                report_ds.push_to_hub(output_dataset, private=private, token=HF_TOKEN)336                break337            except Exception as e:338                logger.error(f"Upload attempt {attempt}/{max_retries} failed: {e}")339                if attempt < max_retries:340                    time.sleep(30 * (2 ** (attempt - 1)))341                else:342                    logger.error("All upload attempts failed.")343                    sys.exit(1)344 345        logger.info(f"Stats pushed to: https://huggingface.co/datasets/{output_dataset}")346 347 348if __name__ == "__main__":349    parser = argparse.ArgumentParser(350        description="Generate statistics for object detection datasets on HF Hub",351        formatter_class=argparse.RawDescriptionHelpFormatter,352        epilog="""353Bbox formats:354  coco_xywh     [x, y, width, height] in pixels (default)355  xyxy          [xmin, ymin, xmax, ymax] in pixels356  voc           [xmin, ymin, xmax, ymax] in pixels (alias for xyxy)357  yolo          [cx, cy, w, h] normalized 0-1358  tfod          [xmin, ymin, xmax, ymax] normalized 0-1359  label_studio  [x, y, w, h] percentage 0-100360 361Examples:362  uv run stats-hf-dataset.py merve/coco-dataset363  uv run stats-hf-dataset.py merve/coco-dataset --top 20 --report json364  uv run stats-hf-dataset.py merve/coco-dataset --streaming --max-samples 5000365        """,366    )367 368    parser.add_argument("input_dataset", help="Input dataset ID on HF Hub")369    parser.add_argument("--bbox-column", default="bbox", help="Column containing bboxes (default: bbox)")370    parser.add_argument("--category-column", default="category", help="Column containing categories (default: category)")371    parser.add_argument("--bbox-format", choices=BBOX_FORMATS, default="coco_xywh", help="Bbox format (default: coco_xywh)")372    parser.add_argument("--width-column", default="width", help="Column for image width (default: width)")373    parser.add_argument("--height-column", default="height", help="Column for image height (default: height)")374    parser.add_argument("--split", default="train", help="Dataset split (default: train)")375    parser.add_argument("--max-samples", type=int, help="Max samples to process")376    parser.add_argument("--streaming", action="store_true", help="Use streaming mode")377    parser.add_argument("--top", type=int, default=10, help="Top-N items for histograms (default: 10)")378    parser.add_argument("--report", choices=["text", "json"], default="text", help="Report format (default: text)")379    parser.add_argument("--tolerance", type=float, default=0.5, help="Out-of-bounds tolerance in pixels (default: 0.5)")380    parser.add_argument("--hf-token", help="HF API token")381    parser.add_argument("--output-dataset", help="Push stats report to this HF dataset")382    parser.add_argument("--private", action="store_true", help="Make output dataset private")383 384    args = parser.parse_args()385 386    main(387        input_dataset=args.input_dataset,388        bbox_column=args.bbox_column,389        category_column=args.category_column,390        bbox_format=args.bbox_format,391        width_column=args.width_column,392        height_column=args.height_column,393        split=args.split,394        max_samples=args.max_samples,395        streaming=args.streaming,396        top=args.top,397        report_format=args.report,398        tolerance=args.tolerance,399        hf_token=args.hf_token,400        output_dataset=args.output_dataset,401        private=args.private,402    )403