import os import gradio as gr import numpy as np import torch import random from PIL import Image, ImageDraw from typing import Iterable from gradio.themes import Soft from gradio.themes.utils import colors, fonts, sizes from transformers import Sam3Processor, Sam3Model # --- Handle optional 'spaces' import for local compatibility --- try: import spaces except ImportError: class spaces: @staticmethod def GPU(duration=60): def decorator(func): return func return decorator colors.steel_blue = colors.Color( name="steel_blue", c50="#EBF3F8", c100="#D3E5F0", c200="#A8CCE1", c300="#7DB3D2", c400="#529AC3", c500="#4682B4", c600="#3E72A0", c700="#36638C", c800="#2E5378", c900="#264364", c950="#1E3450", ) class SteelBlueTheme(Soft): def __init__( self, *, primary_hue: colors.Color | str = colors.gray, secondary_hue: colors.Color | str = colors.steel_blue, neutral_hue: colors.Color | str = colors.slate, text_size: sizes.Size | str = sizes.text_lg, font: fonts.Font | str | Iterable[fonts.Font | str] = ( fonts.GoogleFont("Outfit"), "Arial", "sans-serif", ), font_mono: fonts.Font | str | Iterable[fonts.Font | str] = ( fonts.GoogleFont("IBM Plex Mono"), "ui-monospace", "monospace", ), ): super().__init__( primary_hue=primary_hue, secondary_hue=secondary_hue, neutral_hue=neutral_hue, text_size=text_size, font=font, font_mono=font_mono, ) super().set( background_fill_primary="*primary_50", background_fill_primary_dark="*primary_900", body_background_fill="linear-gradient(135deg, *primary_200, *primary_100)", body_background_fill_dark="linear-gradient(135deg, *primary_900, *primary_800)", button_primary_text_color="white", button_primary_text_color_hover="white", button_primary_background_fill="linear-gradient(90deg, *secondary_500, *secondary_600)", button_primary_background_fill_hover="linear-gradient(90deg, *secondary_600, *secondary_700)", button_primary_background_fill_dark="linear-gradient(90deg, *secondary_600, *secondary_700)", button_primary_background_fill_hover_dark="linear-gradient(90deg, *secondary_500, *secondary_600)", slider_color="*secondary_500", slider_color_dark="*secondary_600", block_title_text_weight="600", block_border_width="3px", block_shadow="*shadow_drop_lg", button_primary_shadow="*shadow_drop_lg", button_large_padding="11px", color_accent_soft="*primary_100", block_label_background_fill="*primary_200", ) steel_blue_theme = SteelBlueTheme() # --- Hardware Setup --- device = "cuda" if torch.cuda.is_available() else "cpu" print(f"Using device: {device}") # --- Model Loading --- # Using the facebook/sam3 model as requested try: print("Loading SAM3 Model and Processor...") model = Sam3Model.from_pretrained("facebook/sam3").to(device) processor = Sam3Processor.from_pretrained("facebook/sam3") print("Model loaded successfully.") except Exception as e: print(f"Error loading model: {e}") print("Ensure you have the correct libraries installed and access to the model.") # Fallback/Placeholder for demonstration if model doesn't exist in environment yet model = None processor = None @spaces.GPU(duration=60) def segment_image(input_image, text_prompt, threshold=0.5): if input_image is None: raise gr.Error("Please upload an image.") if not text_prompt: raise gr.Error("Please enter a text prompt (e.g., 'cat', 'face').") if model is None or processor is None: raise gr.Error("Model not loaded correctly.") # Convert image to RGB image_pil = input_image.convert("RGB") # Preprocess inputs = processor(images=image_pil, text=text_prompt, return_tensors="pt").to(device) # Inference with torch.no_grad(): outputs = model(**inputs) # Post-process results results = processor.post_process_instance_segmentation( outputs, threshold=threshold, mask_threshold=0.5, target_sizes=inputs.get("original_sizes").tolist() )[0] masks = results['masks'] # Boolean tensor [N, H, W] scores = results['scores'] # Prepare for Gradio AnnotatedImage # Gradio expects (image, [(mask, label), ...]) annotations = [] masks_np = masks.cpu().numpy() scores_np = scores.cpu().numpy() for i, mask in enumerate(masks_np): # mask is a boolean array (True/False). # AnnotatedImage handles the coloring automatically. # We just pass the mask and a label. score_val = scores_np[i] label = f"{text_prompt} ({score_val:.2f})" annotations.append((mask, label)) # Return tuple format for AnnotatedImage return (image_pil, annotations) css=""" #col-container { margin: 0 auto; max-width: 980px; } #main-title h1 {font-size: 2.1em !important;} """ with gr.Blocks(css=css, theme=steel_blue_theme) as demo: with gr.Column(elem_id="col-container"): gr.Markdown( "# **SAM3 Image Segmentation**", elem_id="main-title" ) gr.Markdown("Segment objects in images using **SAM3** (Segment Anything Model 3) with text prompts.") with gr.Row(): # Left Column: Inputs with gr.Column(scale=1): input_image = gr.Image(label="Input Image", type="pil", height=300) text_prompt = gr.Textbox( label="Text Prompt", placeholder="e.g., cat, ear, car wheel...", info="What do you want to segment?" ) run_button = gr.Button("Segment", variant="primary") # Right Column: Output with gr.Column(scale=1.5): # AnnotatedImage creates a nice overlay visualization output_image = gr.AnnotatedImage(label="Segmented Output", height=400) with gr.Row(): threshold = gr.Slider(label="Confidence Threshold", minimum=0.0, maximum=1.0, value=0.4, step=0.05) # Examples gr.Examples( examples=[ ["examples/cat.jpg", "cat", 0.5], ["examples/car.jpg", "tire", 0.4], ["examples/fruit.jpg", "apple", 0.5], ], inputs=[input_image, text_prompt, threshold], outputs=[output_image], fn=segment_image, cache_examples=False, label="Examples" ) run_button.click( fn=segment_image, inputs=[input_image, text_prompt, threshold], outputs=[output_image] ) if __name__ == "__main__": demo.launch(ssr_mode=False, show_error=True)