import os import cv2 import tempfile import spaces import gradio as gr import numpy as np import torch import matplotlib import matplotlib.pyplot as plt from PIL import Image from typing import Iterable from gradio.themes import Soft from gradio.themes.utils import colors, fonts, sizes from transformers import ( Sam3Model, Sam3Processor, Sam3VideoModel, Sam3VideoProcessor ) # --- THEME CONFIGURATION --- 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 CustomBlueTheme(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", ) app_theme = CustomBlueTheme() # --- GLOBAL MODEL LOADING --- device = "cuda" if torch.cuda.is_available() else "cpu" print(f"🖥️ Using compute device: {device}") print("⏳ Loading SAM3 Models permanently into memory...") try: # 1. Load Image Segmentation Model print(" ... Loading Image Model") IMG_MODEL = Sam3Model.from_pretrained("facebook/sam3").to(device) IMG_PROCESSOR = Sam3Processor.from_pretrained("facebook/sam3") # 2. Load Video Segmentation Model # Using bfloat16 for video to optimize VRAM usage while keeping speed print(" ... Loading Video Model") VID_MODEL = Sam3VideoModel.from_pretrained("facebook/sam3").to(device, dtype=torch.bfloat16) VID_PROCESSOR = Sam3VideoProcessor.from_pretrained("facebook/sam3") print("✅ All Models loaded successfully!") except Exception as e: print(f"❌ CRITICAL ERROR LOADING MODELS: {e}") IMG_MODEL = None VID_MODEL = None IMG_PROCESSOR = None VID_PROCESSOR = None # --- UTILS --- def apply_mask_overlay(base_image, mask_data, opacity=0.5): """Draws segmentation masks on top of an image.""" if isinstance(base_image, np.ndarray): base_image = Image.fromarray(base_image) base_image = base_image.convert("RGBA") if mask_data is None or len(mask_data) == 0: return base_image.convert("RGB") if isinstance(mask_data, torch.Tensor): mask_data = mask_data.cpu().numpy() mask_data = mask_data.astype(np.uint8) # Handle dimensions if mask_data.ndim == 4: mask_data = mask_data[0] if mask_data.ndim == 3 and mask_data.shape[0] == 1: mask_data = mask_data[0] num_masks = mask_data.shape[0] if mask_data.ndim == 3 else 1 if mask_data.ndim == 2: mask_data = [mask_data] num_masks = 1 try: color_map = matplotlib.colormaps["rainbow"].resampled(max(num_masks, 1)) except AttributeError: import matplotlib.cm as cm color_map = cm.get_cmap("rainbow").resampled(max(num_masks, 1)) rgb_colors = [tuple(int(c * 255) for c in color_map(i)[:3]) for i in range(num_masks)] composite_layer = Image.new("RGBA", base_image.size, (0, 0, 0, 0)) for i, single_mask in enumerate(mask_data): mask_bitmap = Image.fromarray((single_mask * 255).astype(np.uint8)) if mask_bitmap.size != base_image.size: mask_bitmap = mask_bitmap.resize(base_image.size, resample=Image.NEAREST) fill_color = rgb_colors[i] color_fill = Image.new("RGBA", base_image.size, fill_color + (0,)) mask_alpha = mask_bitmap.point(lambda v: int(v * opacity) if v > 0 else 0) color_fill.putalpha(mask_alpha) composite_layer = Image.alpha_composite(composite_layer, color_fill) return Image.alpha_composite(base_image, composite_layer).convert("RGB") # --- GPU INFERENCE FUNCTIONS --- @spaces.GPU def run_image_segmentation(source_img, text_query, conf_thresh=0.5): if IMG_MODEL is None or IMG_PROCESSOR is None: raise gr.Error("Models failed to load on startup. Check logs.") if source_img is None or not text_query: raise gr.Error("Please provide an image and a text prompt.") try: pil_image = source_img.convert("RGB") # Models are already on device, just move inputs model_inputs = IMG_PROCESSOR(images=pil_image, text=text_query, return_tensors="pt").to(device) with torch.no_grad(): inference_output = IMG_MODEL(**model_inputs) processed_results = IMG_PROCESSOR.post_process_instance_segmentation( inference_output, threshold=conf_thresh, mask_threshold=0.5, target_sizes=model_inputs.get("original_sizes").tolist() )[0] # Use AnnotatedImage format annotation_list = [] raw_masks = processed_results['masks'].cpu().numpy() raw_scores = processed_results['scores'].cpu().numpy() for idx, mask_array in enumerate(raw_masks): label_str = f"{text_query} ({raw_scores[idx]:.2f})" annotation_list.append((mask_array, label_str)) return (pil_image, annotation_list) except Exception as e: raise gr.Error(f"Error during image processing: {e}") def calc_timeout_duration(vid_file, *args): return args[-1] if args else 60 @spaces.GPU(duration=calc_timeout_duration) def run_video_segmentation(source_vid, text_query, frame_limit, time_limit): if VID_MODEL is None or VID_PROCESSOR is None: raise gr.Error("Video Models failed to load on startup.") if not source_vid or not text_query: raise gr.Error("Missing video or prompt.") try: video_cap = cv2.VideoCapture(source_vid) vid_fps = video_cap.get(cv2.CAP_PROP_FPS) vid_w = int(video_cap.get(cv2.CAP_PROP_FRAME_WIDTH)) vid_h = int(video_cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) video_frames = [] counter = 0 while video_cap.isOpened(): ret, frame = video_cap.read() if not ret or (frame_limit > 0 and counter >= frame_limit): break video_frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) counter += 1 video_cap.release() # VID_MODEL is already on device in bfloat16 session = VID_PROCESSOR.init_video_session(video=video_frames, inference_device=device, dtype=torch.bfloat16) session = VID_PROCESSOR.add_text_prompt(inference_session=session, text=text_query) temp_out_path = tempfile.mktemp(suffix=".mp4") video_writer = cv2.VideoWriter(temp_out_path, cv2.VideoWriter_fourcc(*'mp4v'), vid_fps, (vid_w, vid_h)) for model_out in VID_MODEL.propagate_in_video_iterator(inference_session=session, max_frame_num_to_track=len(video_frames)): post_processed = VID_PROCESSOR.postprocess_outputs(session, model_out) f_idx = model_out.frame_idx original_pil = Image.fromarray(video_frames[f_idx]) if 'masks' in post_processed: detected_masks = post_processed['masks'] if detected_masks.ndim == 4: detected_masks = detected_masks.squeeze(1) final_frame = apply_mask_overlay(original_pil, detected_masks) else: final_frame = original_pil video_writer.write(cv2.cvtColor(np.array(final_frame), cv2.COLOR_RGB2BGR)) video_writer.release() return temp_out_path, "Video processing completed successfully.✅" except Exception as e: return None, f"Error during video processing: {str(e)}" # --- GUI --- custom_css=""" #col-container { margin: 0 auto; max-width: 1100px; } #main-title h1 { font-size: 2.1em !important; } """ with gr.Blocks(css=custom_css, theme=app_theme) as main_interface: with gr.Column(elem_id="col-container"): gr.Markdown("# **SAM3: Segment Anything Model 3**", elem_id="main-title") gr.Markdown("Segment objects in image or video using **SAM3** (Segment Anything Model 3) with text prompts.") with gr.Tabs(): with gr.Tab("Image Segmentation"): with gr.Row(): with gr.Column(scale=1): image_input = gr.Image(label="Source Image", type="pil", height=350) txt_prompt_img = gr.Textbox(label="Text Description", placeholder="e.g., cat, face, car wheel") with gr.Accordion("Advanced Settings", open=False): conf_slider = gr.Slider(0.0, 1.0, value=0.45, step=0.05, label="Confidence Threshold") btn_process_img = gr.Button("Segment Image", variant="primary") with gr.Column(scale=1.5): image_result = gr.AnnotatedImage(label="Segmented Result", height=410) gr.Examples( examples=[ ["examples/player.jpg", "player in white", 0.5], ], inputs=[image_input, txt_prompt_img, conf_slider], outputs=[image_result], fn=run_image_segmentation, cache_examples=False, label="Image Examples" ) btn_process_img.click( fn=run_image_segmentation, inputs=[image_input, txt_prompt_img, conf_slider], outputs=[image_result] ) with gr.Tab("Video Segmentation"): with gr.Row(): with gr.Column(): video_input = gr.Video(label="Source Video", format="mp4", height=320) txt_prompt_vid = gr.Textbox(label="Text Description", placeholder="e.g., person running, red car") with gr.Row(): frame_limiter = gr.Slider(10, 500, value=60, step=10, label="Max Frames") time_limiter = gr.Radio([60, 120, 180], value=60, label="Timeout (seconds)") btn_process_vid = gr.Button("Segment Video", variant="primary") with gr.Column(): video_result = gr.Video(label="Processed Video") process_status = gr.Textbox(label="System Status", interactive=False) gr.Examples( examples=[ ["examples/sample_video.mp4", "players", 120, 120], ], inputs=[video_input, txt_prompt_vid, frame_limiter, time_limiter], outputs=[video_result, process_status], fn=run_video_segmentation, cache_examples=False, label="Video Examples" ) btn_process_vid.click( run_video_segmentation, inputs=[video_input, txt_prompt_vid, frame_limiter, time_limiter], outputs=[video_result, process_status] ) if __name__ == "__main__": main_interface.launch(ssr_mode=False, mcp_server=True, show_error=True)