uv-scripts/sam3

1#!/usr/bin/env python32# /// script3# requires-python = ">=3.10"4# dependencies = [5#     "transformers@git+https://github.com/huggingface/transformers.git@1fba72361e8e0e865d569f7cd15e5aa50b41ac9a",6#     "datasets",7#     "huggingface-hub",8#     "pillow",9#     "tqdm",10#     "torchvision",11#     "accelerate",12# ]13# ///14 15"""16Detect objects in images using Meta's SAM3 (Segment Anything Model 3).17 18This script processes images from a HuggingFace dataset and detects a single object19type based on a text prompt, outputting bounding boxes in HuggingFace object detection format.20 21Examples:22    # Detect photographs in historical newspapers23    uv run detect-objects.py \\24        davanstrien/newspapers-with-images-after-photography \\25        my-username/newspapers-detected \\26        --class-name photograph27 28    # Detect animals in camera trap images29    uv run detect-objects.py \\30        wildlife-images \\31        wildlife-detected \\32        --class-name animal \\33        --confidence-threshold 0.634 35    # Test on small subset36    uv run detect-objects.py input output \\37        --class-name table \\38        --max-samples 1039 40    # Run on HF Jobs with GPU41    hf jobs uv run --flavor a100-large \\42        -s HF_TOKEN=HF_TOKEN \\43        https://huggingface.co/datasets/uv-scripts/sam3/raw/main/detect-objects.py \\44        input-dataset output-dataset \\45        --class-name photograph \\46        --confidence-threshold 0.547 48Note: To detect multiple object types, run the script multiple times with different49      --class-name values and merge the results.50"""51 52import argparse53import logging54import os55import sys56import time57from typing import Any, Dict, List58 59import torch60from datasets import ClassLabel, Dataset, Features, Sequence, Value, load_dataset61from datasets import Image as ImageFeature62from huggingface_hub import DatasetCard, HfApi, login63from PIL import Image64from tqdm.auto import tqdm65from transformers import Sam3Model, Sam3Processor66 67os.environ["HF_XET_HIGH_PERFORMANCE"] = "1"68# Configure logging69logging.basicConfig(70    level=logging.INFO,71    format="%(asctime)s - %(levelname)s - %(message)s",72    datefmt="%H:%M:%S",73)74logger = logging.getLogger(__name__)75 76# GPU availability check77if not torch.cuda.is_available():78    logger.error("❌ CUDA is not available. This script requires a GPU.")79    logger.error("For local testing, ensure you have a CUDA-capable GPU.")80    logger.error("For cloud execution, use HF Jobs with --flavor l4x1 or similar.")81    sys.exit(1)82 83 84def parse_args():85    """Parse command line arguments."""86    parser = argparse.ArgumentParser(87        description="Detect objects in images using SAM3",88        formatter_class=argparse.RawDescriptionHelpFormatter,89        epilog=__doc__,90    )91 92    # Required arguments93    parser.add_argument(94        "input_dataset", help="Input HuggingFace dataset ID (e.g., 'username/dataset')"95    )96    parser.add_argument(97        "output_dataset", help="Output HuggingFace dataset ID (e.g., 'username/output')"98    )99 100    # Object detection configuration101    parser.add_argument(102        "--class-name",103        required=True,104        help="Object class to detect (e.g., 'photograph', 'animal', 'table')",105    )106    parser.add_argument(107        "--confidence-threshold",108        type=float,109        default=0.5,110        help="Minimum confidence score for detections (default: 0.5)",111    )112    parser.add_argument(113        "--mask-threshold",114        type=float,115        default=0.5,116        help="Threshold for mask generation (default: 0.5)",117    )118 119    # Dataset configuration120    parser.add_argument(121        "--image-column",122        default="image",123        help="Name of the column containing images (default: 'image')",124    )125    parser.add_argument(126        "--split", default="train", help="Dataset split to process (default: 'train')"127    )128    parser.add_argument(129        "--max-samples",130        type=int,131        default=None,132        help="Maximum number of samples to process (for testing)",133    )134    parser.add_argument(135        "--shuffle", action="store_true", help="Shuffle dataset before processing"136    )137 138    # Processing configuration139    parser.add_argument(140        "--batch-size",141        type=int,142        default=4,143        help="Batch size for processing (default: 4)",144    )145    parser.add_argument(146        "--model",147        default="facebook/sam3",148        help="SAM3 model ID (default: 'facebook/sam3')",149    )150    parser.add_argument(151        "--dtype",152        default="bfloat16",153        choices=["float32", "float16", "bfloat16"],154        help="Model precision (default: 'bfloat16')",155    )156 157    # Output configuration158    parser.add_argument(159        "--private", action="store_true", help="Make output dataset private"160    )161    parser.add_argument(162        "--hf-token",163        default=None,164        help="HuggingFace token (default: uses HF_TOKEN env var or cached token)",165    )166 167    return parser.parse_args()168 169 170def create_dataset_card(171    source_dataset: str,172    model: str,173    class_name: str,174    num_samples: int,175    total_detections: int,176    images_with_detections: int,177    processing_time: str,178    confidence_threshold: float,179    mask_threshold: float,180    batch_size: int,181    dtype: str,182    image_column: str = "image",183    split: str = "train",184) -> str:185    """Create a dataset card documenting the object detection process."""186    from datetime import datetime187 188    model_name = model.split("/")[-1]189    avg_detections = total_detections / num_samples if num_samples > 0 else 0190    detection_rate = (191        (images_with_detections / num_samples * 100) if num_samples > 0 else 0192    )193 194    return f"""---195tags:196- object-detection197- sam3198- segment-anything199- bounding-boxes200- uv-script201- generated202---203 204# Object Detection: {class_name.title()} Detection using {model_name}205 206This dataset contains object detection results (bounding boxes) for **{class_name}** detected in images from [{source_dataset}](https://huggingface.co/datasets/{source_dataset}) using Meta's SAM3 (Segment Anything Model 3).207 208**Generated using**: [uv-scripts/sam3](https://huggingface.co/datasets/uv-scripts/sam3) detection script209 210## Detection Statistics211 212- **Objects Detected**: {class_name}213- **Total Detections**: {total_detections:,}214- **Images with Detections**: {images_with_detections:,} / {num_samples:,} ({detection_rate:.1f}%)215- **Average Detections per Image**: {avg_detections:.2f}216 217## Processing Details218 219- **Source Dataset**: [{source_dataset}](https://huggingface.co/datasets/{source_dataset})220- **Model**: [{model}](https://huggingface.co/{model})221- **Script Repository**: [uv-scripts/sam3](https://huggingface.co/datasets/uv-scripts/sam3)222- **Number of Samples Processed**: {num_samples:,}223- **Processing Time**: {processing_time}224- **Processing Date**: {datetime.now().strftime("%Y-%m-%d %H:%M UTC")}225 226### Configuration227 228- **Image Column**: `{image_column}`229- **Dataset Split**: `{split}`230- **Class Name**: `{class_name}`231- **Confidence Threshold**: {confidence_threshold}232- **Mask Threshold**: {mask_threshold}233- **Batch Size**: {batch_size}234- **Model Dtype**: {dtype}235 236## Model Information237 238SAM3 (Segment Anything Model 3) is Meta's state-of-the-art object detection and segmentation model that excels at:239- 🎯 **Zero-shot detection** - Detect objects using natural language prompts240- 📦 **Bounding boxes** - Accurate object localization241- 🎭 **Instance segmentation** - Pixel-perfect masks (not included in this dataset)242- 🖼️ **Any image domain** - Works on photos, documents, medical images, etc.243 244This dataset uses SAM3 in text-prompted detection mode to find instances of "{class_name}" in the source images.245 246## Dataset Structure247 248The dataset contains all original columns from the source dataset plus an `objects` column with detection results in HuggingFace object detection format (dict-of-lists):249 250- **bbox**: List of bounding boxes in `[x, y, width, height]` format (pixel coordinates)251- **category**: List of category indices (always `0` for single-class detection)252- **score**: List of confidence scores (0.0 to 1.0)253 254### Schema255 256```python257{{258    "objects": {{259        "bbox": [[x, y, w, h], ...],      # List of bounding boxes260        "category": [0, 0, ...],           # All same class261        "score": [0.95, 0.87, ...]        # Confidence scores262    }}263}}264```265 266## Usage267 268```python269from datasets import load_dataset270 271# Load the dataset272dataset = load_dataset("{{{{output_dataset_id}}}}", split="{split}")273 274# Access detections for an image275example = dataset[0]276detections = example["objects"]277 278# Iterate through all detected objects in this image279for bbox, category, score in zip(280    detections["bbox"],281    detections["category"],282    detections["score"]283):284    x, y, w, h = bbox285    print(f"Detected {class_name} at ({{x}}, {{y}}) with confidence {{score:.2f}}")286 287# Filter high-confidence detections288high_conf_examples = [289    ex for ex in dataset290    if any(score > 0.8 for score in ex["objects"]["score"])291]292 293# Count total detections across dataset294total = sum(len(ex["objects"]["bbox"]) for ex in dataset)295print(f"Total detections: {{total}}")296```297 298## Visualization299 300To visualize the detections, you can use the visualization script from the same repository:301 302```bash303# Visualize first sample with detections304uv run https://huggingface.co/datasets/uv-scripts/sam3/raw/main/visualize-detections.py \\305    {{{{output_dataset_id}}}} \\306    --first-with-detections307 308# Visualize random samples309uv run https://huggingface.co/datasets/uv-scripts/sam3/raw/main/visualize-detections.py \\310    {{{{output_dataset_id}}}} \\311    --num-samples 5312 313# Save visualizations to files314uv run https://huggingface.co/datasets/uv-scripts/sam3/raw/main/visualize-detections.py \\315    {{{{output_dataset_id}}}} \\316    --num-samples 3 \\317    --output-dir ./visualizations318```319 320## Reproduction321 322This dataset was generated using the [uv-scripts/sam3](https://huggingface.co/datasets/uv-scripts/sam3) object detection script:323 324```bash325uv run https://huggingface.co/datasets/uv-scripts/sam3/raw/main/detect-objects.py \\326    {source_dataset} \\327    <output-dataset> \\328    --class-name {class_name} \\329    --confidence-threshold {confidence_threshold} \\330    --mask-threshold {mask_threshold} \\331    --batch-size {batch_size} \\332    --dtype {dtype}333```334 335### Running on HuggingFace Jobs (GPU)336 337This script requires a GPU. To run on HuggingFace infrastructure:338 339```bash340hf jobs uv run --flavor a100-large \\341    -s HF_TOKEN=HF_TOKEN \\342    https://huggingface.co/datasets/uv-scripts/sam3/raw/main/detect-objects.py \\343    {source_dataset} \\344    <output-dataset> \\345    --class-name {class_name} \\346    --confidence-threshold {confidence_threshold}347```348 349## Performance350 351- **Processing Speed**: ~{num_samples / (float(processing_time.split()[0]) * 60) if processing_time.split()[0].replace(".", "").isdigit() else "N/A":.1f} images/second352- **GPU Configuration**: CUDA with {dtype} precision353 354---355 356Generated with 🤖 [UV Scripts](https://huggingface.co/uv-scripts)357"""358 359 360def load_and_validate_dataset(361    dataset_id: str,362    split: str,363    image_column: str,364    max_samples: int = None,365    shuffle: bool = False,366    hf_token: str = None,367) -> Dataset:368    """Load dataset and validate it has the required image column."""369    logger.info(f"📂 Loading dataset: {dataset_id} (split: {split})")370 371    try:372        dataset = load_dataset(dataset_id, split=split, token=hf_token)373    except Exception as e:374        logger.error(f"Failed to load dataset '{dataset_id}': {e}")375        sys.exit(1)376 377    # Validate image column exists378    if image_column not in dataset.column_names:379        logger.error(f"Column '{image_column}' not found in dataset")380        logger.error(f"Available columns: {dataset.column_names}")381        sys.exit(1)382 383    # Shuffle if requested384    if shuffle:385        logger.info("🔀 Shuffling dataset")386        dataset = dataset.shuffle()387 388    # Limit samples if requested389    if max_samples is not None:390        logger.info(f"🔢 Limiting to {max_samples} samples")391        dataset = dataset.select(range(min(max_samples, len(dataset))))392 393    logger.info(f"✅ Loaded {len(dataset)} samples")394    return dataset395 396 397def process_batch(398    batch: Dict[str, List[Any]],399    image_column: str,400    class_name: str,401    processor: Sam3Processor,402    model: Sam3Model,403    confidence_threshold: float,404    mask_threshold: float,405) -> Dict[str, List[List[Dict[str, Any]]]]:406    """Process a batch of images and return detections for a single class."""407    images = batch[image_column]408 409    # Convert to PIL Images and ensure RGB410    pil_images = []411    for img in images:412        if isinstance(img, str):413            img = Image.open(img)414        if img.mode == "L" or img.mode != "RGB":415            img = img.convert("RGB")416        pil_images.append(img)417 418    # Process batch through model419    try:420        inputs = processor(421            images=pil_images,422            text=[class_name] * len(pil_images),  # Same prompt for all images423            return_tensors="pt",424        ).to(model.device, dtype=model.dtype)425 426        with torch.no_grad():427            outputs = model(**inputs)428 429        # Post-process outputs using original_sizes from processor430        results = processor.post_process_instance_segmentation(431            outputs,432            threshold=confidence_threshold,433            mask_threshold=mask_threshold,434            target_sizes=inputs.get("original_sizes").tolist(),435        )436 437    except Exception as e:438        logger.warning(f"⚠️  Failed to process batch: {e}")439        # Return empty detections for all images in batch440        return {441            "objects": [442                {"bbox": [], "category": [], "score": []}443                for _ in range(len(pil_images))444            ]445        }446 447    # Convert to HuggingFace object detection format (dict-of-lists per image)448    batch_objects = []449 450    for result in results:451        boxes = result.get("boxes", torch.tensor([]))452        scores = result.get("scores", torch.tensor([]))453 454        # Handle empty results455        if len(boxes) == 0:456            batch_objects.append({"bbox": [], "category": [], "score": []})457            continue458 459        # Build lists for this image460        image_bboxes = []461        image_categories = []462        image_scores = []463 464        # Convert to float32 before numpy (bfloat16 not supported by numpy)465        boxes_np = boxes.cpu().float().numpy()466        scores_np = scores.cpu().float().numpy()467 468        for box, score in zip(boxes_np, scores_np):469            x1, y1, x2, y2 = box470            width = x2 - x1471            height = y2 - y1472 473            image_bboxes.append([float(x1), float(y1), float(width), float(height)])474            image_categories.append(0)  # Single class, always index 0475            image_scores.append(float(score))476 477        batch_objects.append(478            {479                "bbox": image_bboxes,480                "category": image_categories,481                "score": image_scores,482            }483        )484 485    return {"objects": batch_objects}486 487 488def main():489    args = parse_args()490 491    class_name = args.class_name.strip()492    if not class_name:493        logger.error("❌ Invalid --class-name argument. Provide a class name.")494        sys.exit(1)495 496    logger.info("🚀 SAM3 Object Detection")497    logger.info(f"   Input: {args.input_dataset}")498    logger.info(f"   Output: {args.output_dataset}")499    logger.info(f"   Class: {class_name}")500    logger.info(f"   Confidence threshold: {args.confidence_threshold}")501    logger.info(f"   Batch size: {args.batch_size}")502 503    # Authentication504    if args.hf_token:505        login(token=args.hf_token)506    elif os.getenv("HF_TOKEN"):507        login(token=os.getenv("HF_TOKEN"))508 509    # Load dataset510    dataset = load_and_validate_dataset(511        args.input_dataset,512        args.split,513        args.image_column,514        args.max_samples,515        args.shuffle,516        args.hf_token,517    )518 519    # Load model520    logger.info(f"🤖 Loading SAM3 model: {args.model}")521    try:522        processor = Sam3Processor.from_pretrained(args.model)523        model = Sam3Model.from_pretrained(524            args.model, torch_dtype=getattr(torch, args.dtype), device_map="auto"525        )526        logger.info(f"✅ Model loaded on {model.device}")527    except Exception as e:528        logger.error(f"❌ Failed to load model: {e}")529        logger.error("Ensure the model exists and you have access permissions")530        sys.exit(1)531 532    # Define output schema before processing (dict-of-lists format for object detection)533    logger.info("📊 Creating output schema...")534    new_features = dataset.features.copy()535    new_features["objects"] = {536        "bbox": Sequence(Sequence(Value("float32"), length=4)),537        "category": Sequence(ClassLabel(names=[class_name])),538        "score": Sequence(Value("float32")),539    }540 541    # Process dataset with explicit output features542    logger.info("🔍 Processing images...")543    start_time = time.time()544    processed_dataset = dataset.map(545        lambda batch: process_batch(546            batch,547            args.image_column,548            class_name,549            processor,550            model,551            args.confidence_threshold,552            args.mask_threshold,553        ),554        batched=True,555        batch_size=args.batch_size,556        features=new_features,557        desc="Detecting objects",558    )559    end_time = time.time()560    processing_time_seconds = end_time - start_time561    processing_time_str = f"{processing_time_seconds / 60:.1f} minutes"562 563    # Calculate statistics564    total_detections = sum(len(objs) for objs in processed_dataset["objects"])565    images_with_detections = sum(len(objs) > 0 for objs in processed_dataset["objects"])566 567    logger.info("✅ Detection complete!")568    logger.info(f"   Total detections: {total_detections}")569    logger.info(570        f"   Images with detections: {images_with_detections}/{len(processed_dataset)}"571    )572    logger.info(573        f"   Average detections per image: {total_detections / len(processed_dataset):.2f}"574    )575 576    # Push to hub577    logger.info(f"📤 Pushing to HuggingFace Hub: {args.output_dataset}")578    try:579        processed_dataset.push_to_hub(args.output_dataset, private=args.private)580        logger.info(581            f"✅ Dataset available at: https://huggingface.co/datasets/{args.output_dataset}"582        )583    except Exception as e:584        logger.error(f"❌ Failed to push to hub: {e}")585        logger.info("💾 Saving locally as backup...")586        processed_dataset.save_to_disk("./output_dataset")587        logger.info("✅ Saved to ./output_dataset")588        sys.exit(1)589 590    # Create and push dataset card591    logger.info("📝 Creating dataset card...")592    card_content = create_dataset_card(593        source_dataset=args.input_dataset,594        model=args.model,595        class_name=class_name,596        num_samples=len(processed_dataset),597        total_detections=total_detections,598        images_with_detections=images_with_detections,599        processing_time=processing_time_str,600        confidence_threshold=args.confidence_threshold,601        mask_threshold=args.mask_threshold,602        batch_size=args.batch_size,603        dtype=args.dtype,604        image_column=args.image_column,605        split=args.split,606    )607 608    card = DatasetCard(card_content)609    card.push_to_hub(args.output_dataset, token=args.hf_token or os.getenv("HF_TOKEN"))610    logger.info("✅ Dataset card created and pushed!")611 612 613if __name__ == "__main__":614    main()615