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"""12Create random or stratified subsets of object detection datasets on HF Hub.13 14Mirrors panlabel's sample command. Supports:15 16- Random sampling: Uniform random selection of N images or a fraction17- Stratified sampling: Category-aware weighted sampling to preserve class distribution18- Category filtering: Select only images containing specific categories19- Category mode: Filter by image-level or annotation-level membership20 21Pushes the resulting subset to a new dataset repo on HF Hub.22 23Examples:24  uv run sample-hf-dataset.py merve/dataset merve/subset -n 50025  uv run sample-hf-dataset.py merve/dataset merve/subset --fraction 0.126  uv run sample-hf-dataset.py merve/dataset merve/subset -n 200 --strategy stratified27  uv run sample-hf-dataset.py merve/dataset merve/subset -n 100 --categories "cat,dog,bird"28"""29 30import argparse31import json32import logging33import os34import random35import sys36import time37from collections import Counter, defaultdict38from datetime import datetime39from typing import Any40 41from datasets import load_dataset42from huggingface_hub import DatasetCard, login43from tqdm.auto import tqdm44 45logging.basicConfig(level=logging.INFO)46logger = logging.getLogger(__name__)47 48 49def get_image_categories(50    example: dict[str, Any],51    category_column: str,52) -> list[str]:53    """Get list of category labels from an example."""54    objects = example.get("objects", example)55    categories = objects.get(category_column, []) or []56    return [str(c) for c in categories if c is not None]57 58 59def create_dataset_card(60    source_dataset: str,61    output_dataset: str,62    strategy: str,63    num_samples: int,64    original_size: int,65    categories_filter: list[str] | None,66    category_mode: str,67    seed: int,68    split: str,69) -> str:70    fraction = num_samples / original_size if original_size > 0 else 071    filter_str = f"\n- **Category Filter**: {', '.join(categories_filter)}" if categories_filter else ""72    return f"""---73tags:74- object-detection75- dataset-subset76- panlabel77- uv-script78- generated79---80 81# Dataset Subset: {strategy} sampling82 83A {strategy} subset of [{source_dataset}](https://huggingface.co/datasets/{source_dataset}).84 85## Details86 87- **Source**: [{source_dataset}](https://huggingface.co/datasets/{source_dataset})88- **Strategy**: {strategy}89- **Samples**: {num_samples:,} / {original_size:,} ({fraction:.1%})90- **Seed**: {seed}91- **Split**: `{split}`92- **Category Mode**: {category_mode}{filter_str}93- **Date**: {datetime.now().strftime("%Y-%m-%d %H:%M UTC")}94 95## Reproduction96 97```bash98uv run sample-hf-dataset.py {source_dataset} {output_dataset} \\99    -n {num_samples} --strategy {strategy} --seed {seed}100```101 102Generated with panlabel-hf (sample-hf-dataset.py)103"""104 105 106def main(107    input_dataset: str,108    output_dataset: str,109    n: int | None = None,110    fraction: float | None = None,111    strategy: str = "random",112    category_column: str = "category",113    categories: list[str] | None = None,114    category_mode: str = "images",115    split: str = "train",116    seed: int = 42,117    hf_token: str | None = None,118    private: bool = False,119    create_pr: bool = False,120):121    """Create a subset of an object detection dataset and push to Hub."""122 123    start_time = datetime.now()124 125    if n is None and fraction is None:126        logger.error("Must specify either -n (count) or --fraction")127        sys.exit(1)128 129    if n is not None and fraction is not None:130        logger.error("Specify only one of -n or --fraction, not both")131        sys.exit(1)132 133    HF_TOKEN = hf_token or os.environ.get("HF_TOKEN")134    if HF_TOKEN:135        login(token=HF_TOKEN)136 137    logger.info(f"Loading dataset: {input_dataset} (split={split})")138    dataset = load_dataset(input_dataset, split=split)139    original_size = len(dataset)140    logger.info(f"Loaded {original_size:,} examples")141 142    # Determine target count143    if fraction is not None:144        target_n = max(1, int(original_size * fraction))145        logger.info(f"Fraction {fraction} -> {target_n:,} samples")146    else:147        target_n = min(n, original_size)148 149    rng = random.Random(seed)150 151    # Category filtering152    if categories:153        logger.info(f"Filtering by categories: {categories} (mode={category_mode})")154        keep_indices = []155        for idx in tqdm(range(original_size), desc="Filtering"):156            ex = dataset[idx]157            img_cats = get_image_categories(ex, category_column)158            if category_mode == "images":159                # Keep image if ANY of its annotations match160                if any(c in categories for c in img_cats):161                    keep_indices.append(idx)162            else:  # annotations mode — just check presence, filtering happens below163                if any(c in categories for c in img_cats):164                    keep_indices.append(idx)165 166        dataset = dataset.select(keep_indices)167        logger.info(f"After category filter: {len(dataset):,} examples")168        target_n = min(target_n, len(dataset))169 170    if strategy == "random":171        logger.info(f"Random sampling {target_n:,} from {len(dataset):,}")172        indices = list(range(len(dataset)))173        rng.shuffle(indices)174        selected = sorted(indices[:target_n])175        dataset = dataset.select(selected)176 177    elif strategy == "stratified":178        logger.info(f"Stratified sampling {target_n:,} from {len(dataset):,}")179 180        # Count categories per image and build index181        cat_to_images = defaultdict(list)182        for idx in tqdm(range(len(dataset)), desc="Indexing categories"):183            ex = dataset[idx]184            img_cats = set(get_image_categories(ex, category_column))185            for cat in img_cats:186                cat_to_images[cat].append(idx)187 188        # Compute per-category allocation proportional to frequency189        total_cat_count = sum(len(imgs) for imgs in cat_to_images.values())190        cat_allocations = {}191        for cat, imgs in cat_to_images.items():192            cat_allocations[cat] = max(1, round(target_n * len(imgs) / total_cat_count))193 194        # Greedy selection: pick from underrepresented categories first195        selected = set()196        cat_fulfilled = Counter()197 198        # Sort categories by allocation (smallest first for better representation)199        sorted_cats = sorted(cat_allocations.keys(), key=lambda c: cat_allocations[c])200 201        for cat in sorted_cats:202            needed = cat_allocations[cat] - cat_fulfilled[cat]203            if needed <= 0:204                continue205 206            available = [i for i in cat_to_images[cat] if i not in selected]207            rng.shuffle(available)208            pick = available[:needed]209            selected.update(pick)210 211            # Update fulfilled counts for all categories of picked images212            for idx in pick:213                ex = dataset[idx]214                for c in set(get_image_categories(ex, category_column)):215                    cat_fulfilled[c] += 1216 217        # If we still need more, fill randomly218        if len(selected) < target_n:219            remaining = [i for i in range(len(dataset)) if i not in selected]220            rng.shuffle(remaining)221            selected.update(remaining[: target_n - len(selected)])222 223        # If we have too many, trim224        selected_list = sorted(selected)225        if len(selected_list) > target_n:226            rng.shuffle(selected_list)227            selected_list = sorted(selected_list[:target_n])228 229        dataset = dataset.select(selected_list)230        logger.info(f"Selected {len(dataset):,} samples via stratified sampling")231 232    else:233        logger.error(f"Unknown strategy: {strategy}")234        sys.exit(1)235 236    num_samples = len(dataset)237    processing_duration = datetime.now() - start_time238    processing_time_str = f"{processing_duration.total_seconds():.1f}s"239 240    # Push to Hub241    logger.info(f"Pushing {num_samples:,} samples to {output_dataset}")242    max_retries = 3243    for attempt in range(1, max_retries + 1):244        try:245            if attempt > 1:246                logger.warning("Disabling XET (fallback to HTTP upload)")247                os.environ["HF_HUB_DISABLE_XET"] = "1"248            dataset.push_to_hub(249                output_dataset,250                private=private,251                token=HF_TOKEN,252                max_shard_size="500MB",253                create_pr=create_pr,254            )255            break256        except Exception as e:257            logger.error(f"Upload attempt {attempt}/{max_retries} failed: {e}")258            if attempt < max_retries:259                delay = 30 * (2 ** (attempt - 1))260                logger.info(f"Retrying in {delay}s...")261                time.sleep(delay)262            else:263                logger.error("All upload attempts failed.")264                sys.exit(1)265 266    # Push dataset card267    card_content = create_dataset_card(268        source_dataset=input_dataset,269        output_dataset=output_dataset,270        strategy=strategy,271        num_samples=num_samples,272        original_size=original_size,273        categories_filter=categories,274        category_mode=category_mode,275        seed=seed,276        split=split,277    )278    card = DatasetCard(card_content)279    card.push_to_hub(output_dataset, token=HF_TOKEN)280 281    logger.info("Done!")282    logger.info(f"Dataset: https://huggingface.co/datasets/{output_dataset}")283    logger.info(f"Sampled {num_samples:,} / {original_size:,} in {processing_time_str}")284 285 286if __name__ == "__main__":287    parser = argparse.ArgumentParser(288        description="Create random or stratified subsets of HF object detection datasets",289        formatter_class=argparse.RawDescriptionHelpFormatter,290        epilog="""291Strategies:292  random       Uniform random selection (default)293  stratified   Category-aware weighted sampling294 295Category modes (with --categories):296  images       Keep images containing any matching annotation (default)297  annotations  Keep images containing any matching annotation298 299Examples:300  uv run sample-hf-dataset.py merve/dataset merve/subset -n 500301  uv run sample-hf-dataset.py merve/dataset merve/subset --fraction 0.1302  uv run sample-hf-dataset.py merve/dataset merve/subset -n 200 --strategy stratified303  uv run sample-hf-dataset.py merve/dataset merve/subset -n 100 --categories "cat,dog"304        """,305    )306 307    parser.add_argument("input_dataset", help="Input dataset ID on HF Hub")308    parser.add_argument("output_dataset", help="Output dataset ID on HF Hub")309    parser.add_argument("-n", type=int, help="Number of samples to select")310    parser.add_argument("--fraction", type=float, help="Fraction of dataset to select (0.0-1.0)")311    parser.add_argument("--strategy", choices=["random", "stratified"], default="random", help="Sampling strategy (default: random)")312    parser.add_argument("--category-column", default="category", help="Column containing categories (default: category)")313    parser.add_argument("--categories", help="Comma-separated list of categories to filter by")314    parser.add_argument("--category-mode", choices=["images", "annotations"], default="images", help="How to apply category filter (default: images)")315    parser.add_argument("--split", default="train", help="Dataset split (default: train)")316    parser.add_argument("--seed", type=int, default=42, help="Random seed (default: 42)")317    parser.add_argument("--hf-token", help="HF API token")318    parser.add_argument("--private", action="store_true", help="Make output dataset private")319    parser.add_argument("--create-pr", action="store_true", help="Create PR instead of direct push")320 321    args = parser.parse_args()322 323    cats = None324    if args.categories:325        cats = [c.strip() for c in args.categories.split(",")]326 327    main(328        input_dataset=args.input_dataset,329        output_dataset=args.output_dataset,330        n=args.n,331        fraction=args.fraction,332        strategy=args.strategy,333        category_column=args.category_column,334        categories=cats,335        category_mode=args.category_mode,336        split=args.split,337        seed=args.seed,338        hf_token=args.hf_token,339        private=args.private,340        create_pr=args.create_pr,341    )342