uv-scripts/sam3

1#!/usr/bin/env python32# /// script3# requires-python = ">=3.10"4# dependencies = [5#     "datasets",6#     "matplotlib",7#     "pillow",8# ]9# ///10 11"""12Visualize object detection predictions from a HuggingFace dataset.13 14This script loads a dataset with object detection predictions and visualizes15the bounding boxes on sample images.16 17Examples:18    # Visualize the first sample with detections19    uv run visualize-detections.py my-username/detected-objects --first-with-detections20 21    # Visualize a specific sample22    uv run visualize-detections.py my-username/detected-objects --index 023 24    # Visualize multiple random samples25    uv run visualize-detections.py my-username/detected-objects --num-samples 526 27    # Save visualizations to files instead of displaying28    uv run visualize-detections.py my-username/detected-objects --num-samples 3 --output-dir ./visualizations29 30    # Visualize specific split31    uv run visualize-detections.py my-username/detected-objects --split train --num-samples 532"""33 34import argparse35import random36from pathlib import Path37 38import matplotlib.patches as patches39import matplotlib.pyplot as plt40from datasets import load_dataset41 42 43def parse_args():44    """Parse command line arguments."""45    parser = argparse.ArgumentParser(46        description="Visualize object detection predictions",47        formatter_class=argparse.RawDescriptionHelpFormatter,48        epilog=__doc__,49    )50 51    parser.add_argument(52        "dataset_id", help="HuggingFace dataset ID (e.g., 'username/dataset')"53    )54    parser.add_argument(55        "--index",56        type=int,57        default=None,58        help="Index of sample to visualize (default: random)",59    )60    parser.add_argument(61        "--num-samples",62        type=int,63        default=1,64        help="Number of samples to visualize (default: 1)",65    )66    parser.add_argument(67        "--first-with-detections",68        action="store_true",69        help="Find and visualize the first sample with detections",70    )71    parser.add_argument(72        "--split", default="train", help="Dataset split to use (default: 'train')"73    )74    parser.add_argument(75        "--image-column",76        default="image",77        help="Name of the image column (default: 'image')",78    )79    parser.add_argument(80        "--objects-column",81        default="objects",82        help="Name of the objects column (default: 'objects')",83    )84    parser.add_argument(85        "--output-dir",86        type=str,87        default=None,88        help="Directory to save visualizations (default: show interactively)",89    )90    parser.add_argument(91        "--figsize-width",92        type=int,93        default=15,94        help="Figure width in inches (default: 15)",95    )96    parser.add_argument(97        "--figsize-height",98        type=int,99        default=20,100        help="Figure height in inches (default: 20)",101    )102    parser.add_argument(103        "--bbox-color",104        default="red",105        help="Color for bounding boxes (default: 'red')",106    )107    parser.add_argument(108        "--show-scores",109        action="store_true",110        default=True,111        help="Show confidence scores on bounding boxes",112    )113 114    return parser.parse_args()115 116 117def visualize_sample(118    sample,119    image_column="image",120    objects_column="objects",121    figsize=(15, 20),122    bbox_color="red",123    show_scores=True,124    title=None,125):126    """Visualize a single sample with bounding boxes."""127    image = sample[image_column]128    objects = sample[objects_column]129 130    fig, ax = plt.subplots(1, figsize=figsize)131    ax.imshow(image, cmap="gray" if image.mode == "L" else None)132 133    # Draw bounding boxes134    num_detections = len(objects["bbox"])135    for i in range(num_detections):136        bbox = objects["bbox"][i]137        score = objects["score"][i]138        category = objects["category"][i]139 140        x, y, w, h = bbox141        rect = patches.Rectangle(142            (x, y), w, h, linewidth=2, edgecolor=bbox_color, facecolor="none"143        )144        ax.add_patch(rect)145 146        if show_scores:147            label = f"{score:.2f}"148            ax.text(149                x,150                y - 5,151                label,152                color=bbox_color,153                fontsize=10,154                bbox=dict(facecolor="white", alpha=0.7),155            )156 157    # Set title158    if title:159        ax.set_title(title, fontsize=14, pad=20)160    else:161        ax.set_title(f"Detections: {num_detections}", fontsize=14, pad=20)162 163    ax.axis("off")164    plt.tight_layout()165 166    return fig, ax167 168 169def main():170    args = parse_args()171 172    # Load dataset173    print(f"📂 Loading dataset: {args.dataset_id} (split: {args.split})")174    dataset = load_dataset(args.dataset_id, split=args.split)175    print(f"✅ Loaded {len(dataset)} samples")176 177    # Determine indices to visualize178    if args.index is not None:179        indices = [args.index]180    elif args.first_with_detections:181        # Find first sample with detections182        print("🔍 Finding first sample with detections...")183        first_idx = None184        for idx in range(len(dataset)):185            sample = dataset[idx]186            if len(sample[args.objects_column]["bbox"]) > 0:187                first_idx = idx188                break189 190        if first_idx is None:191            print("❌ No samples with detections found in dataset")192            return193 194        print(f"✅ Found first sample with detections at index {first_idx}")195        indices = [first_idx]196    else:197        # Select random samples198        indices = random.sample(range(len(dataset)), min(args.num_samples, len(dataset)))199 200    # Create output directory if saving201    if args.output_dir:202        output_path = Path(args.output_dir)203        output_path.mkdir(parents=True, exist_ok=True)204        print(f"💾 Saving visualizations to: {output_path}")205 206    # Visualize samples207    figsize = (args.figsize_width, args.figsize_height)208 209    for idx in indices:210        sample = dataset[idx]211        num_detections = len(sample[args.objects_column]["bbox"])212 213        print(f"\n🖼️  Sample {idx}: {num_detections} detections")214 215        # Create visualization216        title = f"Sample {idx} - {num_detections} detections"217        fig, ax = visualize_sample(218            sample,219            image_column=args.image_column,220            objects_column=args.objects_column,221            figsize=figsize,222            bbox_color=args.bbox_color,223            show_scores=args.show_scores,224            title=title,225        )226 227        # Save or show228        if args.output_dir:229            output_file = output_path / f"sample_{idx}.png"230            plt.savefig(output_file, dpi=150, bbox_inches="tight")231            print(f"   Saved: {output_file}")232            plt.close(fig)233        else:234            plt.show()235 236    if args.output_dir:237        print(f"\n✅ Saved {len(indices)} visualizations to {args.output_dir}")238 239 240if __name__ == "__main__":241    main()242