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"""12Validate object detection annotations in a Hugging Face dataset.13 14Streams a HF dataset and checks for common annotation issues, mirroring15panlabel's validate command. Checks include:16 17- Duplicate image file names18- Missing or empty bounding boxes19- Bounding box ordering (xmin <= xmax, ymin <= ymax)20- Bounding boxes out of image bounds21- Non-finite coordinates (NaN/Inf)22- Zero-area bounding boxes23- Empty or missing category labels24- Category ID consistency25 26Supports COCO-style (xywh), XYXY/VOC, YOLO (normalized center xywh),27TFOD (normalized xyxy), and Label Studio (percentage xywh) bbox formats.28Outputs a validation report as text or JSON.29 30Examples:31  uv run validate-hf-dataset.py merve/test-coco-dataset32  uv run validate-hf-dataset.py merve/test-coco-dataset --bbox-format xyxy --strict33  uv run validate-hf-dataset.py merve/test-coco-dataset --bbox-format tfod --report json34  uv run validate-hf-dataset.py merve/test-coco-dataset --report json --max-samples 100035"""36 37import argparse38import json39import logging40import math41import os42import sys43import time44from collections import Counter, defaultdict45from datetime import datetime46from typing import Any47 48from datasets import load_dataset49from huggingface_hub import DatasetCard, login50from tqdm.auto import tqdm51 52logging.basicConfig(level=logging.INFO)53logger = logging.getLogger(__name__)54 55BBOX_FORMATS = ["coco_xywh", "xyxy", "voc", "yolo", "tfod", "label_studio"]56 57 58def to_xyxy(bbox: list[float], fmt: str, img_w: float = 1.0, img_h: float = 1.0) -> tuple[float, float, float, float]:59    """Convert any bbox format to (xmin, ymin, xmax, ymax) in pixel space."""60    if fmt == "coco_xywh":61        x, y, w, h = bbox62        return (x, y, x + w, y + h)63    elif fmt in ("xyxy", "voc"):64        return tuple(bbox[:4])65    elif fmt == "yolo":66        cx, cy, w, h = bbox67        xmin = (cx - w / 2) * img_w68        ymin = (cy - h / 2) * img_h69        xmax = (cx + w / 2) * img_w70        ymax = (cy + h / 2) * img_h71        return (xmin, ymin, xmax, ymax)72    elif fmt == "tfod":73        xmin_n, ymin_n, xmax_n, ymax_n = bbox74        return (xmin_n * img_w, ymin_n * img_h, xmax_n * img_w, ymax_n * img_h)75    elif fmt == "label_studio":76        x_pct, y_pct, w_pct, h_pct = bbox77        return (78            x_pct / 100.0 * img_w,79            y_pct / 100.0 * img_h,80            (x_pct + w_pct) / 100.0 * img_w,81            (y_pct + h_pct) / 100.0 * img_h,82        )83    else:84        raise ValueError(f"Unknown bbox format: {fmt}")85 86 87def is_finite(val: float) -> bool:88    return not (math.isnan(val) or math.isinf(val))89 90 91def validate_example(92    example: dict[str, Any],93    idx: int,94    bbox_column: str,95    category_column: str,96    bbox_format: str,97    image_column: str,98    width_column: str | None,99    height_column: str | None,100    tolerance: float = 0.5,101) -> list[dict]:102    """Validate a single example. Returns a list of issue dicts."""103    issues = []104 105    def add_issue(level: str, code: str, message: str, ann_idx: int | None = None):106        issue = {"level": level, "code": code, "message": message, "example_idx": idx}107        if ann_idx is not None:108            issue["annotation_idx"] = ann_idx109        issues.append(issue)110 111    # Get objects container — handle nested dict (objects column) or flat lists112    objects = example.get("objects", example)113    bboxes = objects.get(bbox_column, [])114    categories = objects.get(category_column, [])115 116    if bboxes is None:117        bboxes = []118    if categories is None:119        categories = []120 121    # Image dimensions (if available)122    img_w = None123    img_h = None124    if width_column and width_column in example:125        img_w = example[width_column]126    elif width_column and objects and width_column in objects:127        img_w = objects[width_column]128    if height_column and height_column in example:129        img_h = example[height_column]130    elif height_column and objects and height_column in objects:131        img_h = objects[height_column]132 133    if not bboxes and not categories:134        add_issue("warning", "W001", "No annotations found in this example")135        return issues136 137    if len(bboxes) != len(categories):138        add_issue(139            "error",140            "E001",141            f"Bbox count ({len(bboxes)}) != category count ({len(categories)})",142        )143 144    for ann_idx, bbox in enumerate(bboxes):145        if bbox is None or len(bbox) < 4:146            add_issue("error", "E002", f"Invalid bbox (need 4 values, got {bbox})", ann_idx)147            continue148 149        # Check finite150        if not all(is_finite(v) for v in bbox[:4]):151            add_issue("error", "E003", f"Non-finite bbox coordinates: {bbox}", ann_idx)152            continue153 154        # Convert to xyxy155        w_for_conv = img_w if img_w else 1.0156        h_for_conv = img_h if img_h else 1.0157        xmin, ymin, xmax, ymax = to_xyxy(bbox[:4], bbox_format, w_for_conv, h_for_conv)158 159        # Check ordering160        if xmin > xmax:161            add_issue("error", "E004", f"xmin ({xmin}) > xmax ({xmax})", ann_idx)162        if ymin > ymax:163            add_issue("error", "E005", f"ymin ({ymin}) > ymax ({ymax})", ann_idx)164 165        # Check zero area166        area = (xmax - xmin) * (ymax - ymin)167        if area <= 0:168            add_issue("warning", "W002", f"Zero or negative area bbox: {bbox}", ann_idx)169 170        # Check bounds (only if image dimensions available)171        if img_w is not None and img_h is not None:172            if xmin < -tolerance or ymin < -tolerance:173                add_issue(174                    "warning",175                    "W003",176                    f"Bbox extends before image origin: ({xmin}, {ymin})",177                    ann_idx,178                )179            if xmax > img_w + tolerance or ymax > img_h + tolerance:180                add_issue(181                    "warning",182                    "W004",183                    f"Bbox extends beyond image bounds: ({xmax}, {ymax}) > ({img_w}, {img_h})",184                    ann_idx,185                )186 187    # Check categories188    for ann_idx, cat in enumerate(categories):189        if cat is None or (isinstance(cat, str) and cat.strip() == ""):190            add_issue("warning", "W005", "Empty category label", ann_idx)191 192    return issues193 194 195def main(196    input_dataset: str,197    bbox_column: str = "bbox",198    category_column: str = "category",199    bbox_format: str = "coco_xywh",200    image_column: str = "image",201    width_column: str | None = "width",202    height_column: str | None = "height",203    split: str = "train",204    max_samples: int | None = None,205    streaming: bool = False,206    strict: bool = False,207    report_format: str = "text",208    tolerance: float = 0.5,209    hf_token: str | None = None,210    output_dataset: str | None = None,211    private: bool = False,212):213    """Validate an object detection dataset from HF Hub."""214 215    start_time = datetime.now()216 217    HF_TOKEN = hf_token or os.environ.get("HF_TOKEN")218    if HF_TOKEN:219        login(token=HF_TOKEN)220 221    logger.info(f"Loading dataset: {input_dataset} (split={split}, streaming={streaming})")222    dataset = load_dataset(input_dataset, split=split, streaming=streaming)223 224    all_issues = []225    file_names = []226    total_annotations = 0227    total_examples = 0228    category_counts = Counter()229    error_count = 0230    warning_count = 0231 232    iterable = dataset233    if max_samples:234        if streaming:235            iterable = dataset.take(max_samples)236        else:237            iterable = dataset.select(range(min(max_samples, len(dataset))))238 239    for idx, example in enumerate(tqdm(iterable, desc="Validating", total=max_samples)):240        total_examples += 1241 242        issues = validate_example(243            example=example,244            idx=idx,245            bbox_column=bbox_column,246            category_column=category_column,247            bbox_format=bbox_format,248            image_column=image_column,249            width_column=width_column,250            height_column=height_column,251            tolerance=tolerance,252        )253        all_issues.extend(issues)254 255        # Count stats256        objects = example.get("objects", example)257        bboxes = objects.get(bbox_column, []) or []258        categories = objects.get(category_column, []) or []259        total_annotations += len(bboxes)260        for cat in categories:261            if cat is not None:262                category_counts[str(cat)] += 1263 264        # Track file names for duplicate check265        fname = example.get("file_name") or example.get("image_id") or str(idx)266        file_names.append(fname)267 268    # Check duplicate file names269    fname_counts = Counter(file_names)270    duplicates = {k: v for k, v in fname_counts.items() if v > 1}271    for fname, count in duplicates.items():272        all_issues.append({273            "level": "warning",274            "code": "W006",275            "message": f"Duplicate file name '{fname}' appears {count} times",276            "example_idx": None,277        })278 279    for issue in all_issues:280        if issue["level"] == "error":281            error_count += 1282        else:283            warning_count += 1284 285    processing_time = datetime.now() - start_time286 287    # Build report288    report = {289        "dataset": input_dataset,290        "split": split,291        "total_examples": total_examples,292        "total_annotations": total_annotations,293        "unique_categories": len(category_counts),294        "errors": error_count,295        "warnings": warning_count,296        "duplicate_filenames": len(duplicates),297        "issues": all_issues,298        "processing_time_seconds": processing_time.total_seconds(),299        "timestamp": datetime.now().isoformat(),300        "valid": error_count == 0 and (not strict or warning_count == 0),301    }302 303    if report_format == "json":304        print(json.dumps(report, indent=2))305    else:306        print("\n" + "=" * 60)307        print(f"Validation Report: {input_dataset}")308        print("=" * 60)309        print(f"  Examples:      {total_examples:,}")310        print(f"  Annotations:   {total_annotations:,}")311        print(f"  Categories:    {len(category_counts):,}")312        print(f"  Errors:        {error_count}")313        print(f"  Warnings:      {warning_count}")314        if duplicates:315            print(f"  Duplicate IDs: {len(duplicates)}")316        print(f"  Processing:    {processing_time.total_seconds():.1f}s")317        print()318 319        if all_issues:320            print("Issues:")321            # Group by code322            by_code = defaultdict(list)323            for issue in all_issues:324                by_code[issue["code"]].append(issue)325 326            for code in sorted(by_code.keys()):327                code_issues = by_code[code]328                level = code_issues[0]["level"].upper()329                sample = code_issues[0]["message"]330                print(f"  [{level}] {code}: {sample}")331                if len(code_issues) > 1:332                    print(f"         ... and {len(code_issues) - 1} more")333            print()334 335        status = "VALID" if report["valid"] else "INVALID"336        mode = " (strict)" if strict else ""337        print(f"Result: {status}{mode}")338        print("=" * 60)339 340    # Optionally push validation report as a dataset341    if output_dataset:342        from datasets import Dataset as HFDataset343 344        report_ds = HFDataset.from_dict({345            "report": [json.dumps(report)],346            "dataset": [input_dataset],347            "valid": [report["valid"]],348            "errors": [error_count],349            "warnings": [warning_count],350            "total_examples": [total_examples],351            "total_annotations": [total_annotations],352            "timestamp": [datetime.now().isoformat()],353        })354 355        logger.info(f"Pushing validation report to {output_dataset}")356        max_retries = 3357        for attempt in range(1, max_retries + 1):358            try:359                if attempt > 1:360                    os.environ["HF_HUB_DISABLE_XET"] = "1"361                report_ds.push_to_hub(362                    output_dataset,363                    private=private,364                    token=HF_TOKEN,365                )366                break367            except Exception as e:368                logger.error(f"Upload attempt {attempt}/{max_retries} failed: {e}")369                if attempt < max_retries:370                    time.sleep(30 * (2 ** (attempt - 1)))371                else:372                    logger.error("All upload attempts failed.")373                    sys.exit(1)374 375        logger.info(f"Report pushed to: https://huggingface.co/datasets/{output_dataset}")376 377    if not report["valid"]:378        sys.exit(1 if strict else 0)379 380 381if __name__ == "__main__":382    parser = argparse.ArgumentParser(383        description="Validate object detection annotations in a HF dataset",384        formatter_class=argparse.RawDescriptionHelpFormatter,385        epilog="""386Bbox formats:387  coco_xywh     [x, y, width, height] in pixels (default)388  xyxy          [xmin, ymin, xmax, ymax] in pixels389  voc           [xmin, ymin, xmax, ymax] in pixels (alias for xyxy)390  yolo          [cx, cy, w, h] normalized 0-1391  tfod          [xmin, ymin, xmax, ymax] normalized 0-1392  label_studio  [x, y, w, h] percentage 0-100393 394Issue codes:395  E001  Bbox/category count mismatch396  E002  Invalid bbox (missing values)397  E003  Non-finite coordinates (NaN/Inf)398  E004  xmin > xmax399  E005  ymin > ymax400  W001  No annotations in example401  W002  Zero or negative area402  W003  Bbox before image origin403  W004  Bbox beyond image bounds404  W005  Empty category label405  W006  Duplicate file name406 407Examples:408  uv run validate-hf-dataset.py merve/coco-dataset409  uv run validate-hf-dataset.py merve/coco-dataset --bbox-format xyxy --strict410  uv run validate-hf-dataset.py merve/coco-dataset --streaming --max-samples 500411        """,412    )413 414    parser.add_argument("input_dataset", help="Input dataset ID on HF Hub")415    parser.add_argument("--bbox-column", default="bbox", help="Column containing bboxes (default: bbox)")416    parser.add_argument("--category-column", default="category", help="Column containing categories (default: category)")417    parser.add_argument(418        "--bbox-format",419        choices=BBOX_FORMATS,420        default="coco_xywh",421        help="Bounding box format (default: coco_xywh)",422    )423    parser.add_argument("--image-column", default="image", help="Column containing images (default: image)")424    parser.add_argument("--width-column", default="width", help="Column for image width (default: width)")425    parser.add_argument("--height-column", default="height", help="Column for image height (default: height)")426    parser.add_argument("--split", default="train", help="Dataset split (default: train)")427    parser.add_argument("--max-samples", type=int, help="Max samples to validate")428    parser.add_argument("--streaming", action="store_true", help="Use streaming mode (no full download)")429    parser.add_argument("--strict", action="store_true", help="Treat warnings as errors")430    parser.add_argument("--report", choices=["text", "json"], default="text", help="Report format (default: text)")431    parser.add_argument("--tolerance", type=float, default=0.5, help="Out-of-bounds tolerance in pixels (default: 0.5)")432    parser.add_argument("--hf-token", help="HF API token")433    parser.add_argument("--output-dataset", help="Push validation report to this HF dataset")434    parser.add_argument("--private", action="store_true", help="Make output dataset private")435 436    args = parser.parse_args()437 438    main(439        input_dataset=args.input_dataset,440        bbox_column=args.bbox_column,441        category_column=args.category_column,442        bbox_format=args.bbox_format,443        image_column=args.image_column,444        width_column=args.width_column,445        height_column=args.height_column,446        split=args.split,447        max_samples=args.max_samples,448        streaming=args.streaming,449        strict=args.strict,450        report_format=args.report,451        tolerance=args.tolerance,452        hf_token=args.hf_token,453        output_dataset=args.output_dataset,454        private=args.private,455    )456