import gradio as gr import torch from transformers import AutoProcessor, AutoModel from PIL import Image, ImageDraw, ImageFont import numpy as np import random import os import wget import traceback # --- Configuration & Model Loading --- # Device Selection with fallback DEVICE = "cuda" if torch.cuda.is_available() and torch.cuda.current_device() >= 0 else "cpu" print(f"Using device: {DEVICE}") # --- CLIP Setup --- CLIP_MODEL_ID = "openai/clip-vit-base-patch32" clip_processor = None clip_model = None def load_clip_model(): global clip_processor, clip_model if clip_processor is None: try: print(f"Loading CLIP processor: {CLIP_MODEL_ID}...") clip_processor = AutoProcessor.from_pretrained(CLIP_MODEL_ID) print("CLIP processor loaded.") except Exception as e: print(f"Error loading CLIP processor: {e}") return False if clip_model is None: try: print(f"Loading CLIP model: {CLIP_MODEL_ID}...") clip_model = AutoModel.from_pretrained(CLIP_MODEL_ID).to(DEVICE) print(f"CLIP model loaded to {DEVICE}.") except Exception as e: print(f"Error loading CLIP model: {e}") return False return True # --- FastSAM Setup --- FASTSAM_CHECKPOINT = "FastSAM-s.pt" FASTSAM_CHECKPOINT_URL = f"https://huggingface.co/CASIA-IVA-Lab/FastSAM-s/resolve/main/{FASTSAM_CHECKPOINT}" fastsam_model = None fastsam_lib_imported = False def check_and_import_fastsam(): global fastsam_lib_imported if not fastsam_lib_imported: try: from fastsam import FastSAM, FastSAMPrompt globals()['FastSAM'] = FastSAM globals()['FastSAMPrompt'] = FastSAMPrompt fastsam_lib_imported = True print("fastsam library imported successfully.") except ImportError as e: print(f"Error: 'fastsam' library not found. Install with 'pip install fastsam': {e}") fastsam_lib_imported = False except Exception as e: print(f"Unexpected error during fastsam import: {e}") traceback.print_exc() fastsam_lib_imported = False return fastsam_lib_imported def download_fastsam_weights(retries=3): if not os.path.exists(FASTSAM_CHECKPOINT): print(f"Downloading FastSAM weights: {FASTSAM_CHECKPOINT} from {FASTSAM_CHECKPOINT_URL}...") for attempt in range(retries): try: wget.download(FASTSAM_CHECKPOINT_URL, FASTSAM_CHECKPOINT) print("FastSAM weights downloaded.") break except Exception as e: print(f"Attempt {attempt + 1}/{retries} failed: {e}") if attempt + 1 == retries: print("Failed to download weights after all attempts.") return False return os.path.exists(FASTSAM_CHECKPOINT) def load_fastsam_model(): global fastsam_model if fastsam_model is None: if not check_and_import_fastsam(): print("Cannot load FastSAM model due to library import failure.") return False if download_fastsam_weights(): try: print(f"Loading FastSAM model: {FASTSAM_CHECKPOINT}...") fastsam_model = FastSAM(FASTSAM_CHECKPOINT) print("FastSAM model loaded.") return True except Exception as e: print(f"Error loading FastSAM model: {e}") traceback.print_exc() return False else: print("FastSAM weights not found or download failed.") return False return True # --- Processing Functions --- def run_clip_zero_shot(image: Image.Image, text_labels: str): if clip_model is None or clip_processor is None: if not load_clip_model(): return "Error: CLIP Model could not be loaded.", None if image is None: return "Please upload an image.", None if not text_labels: return {}, image labels = [label.strip() for label in text_labels.split(',') if label.strip()] if not labels: return {}, image print(f"Running CLIP zero-shot classification with labels: {labels}") try: if image.mode != "RGB": image = image.convert("RGB") inputs = clip_processor(text=labels, images=image, return_tensors="pt", padding=True).to(DEVICE) with torch.no_grad(): outputs = clip_model(**inputs) probs = outputs.logits_per_image.softmax(dim=1) confidences = {labels[i]: float(probs[0, i].item()) for i in range(len(labels))} return confidences, image except Exception as e: print(f"Error during CLIP processing: {e}") traceback.print_exc() return f"Error: {e}", image def run_fastsam_segmentation(image_pil: Image.Image, conf_threshold: float = 0.4, iou_threshold: float = 0.9): if not load_fastsam_model() or not fastsam_lib_imported: return "Error: FastSAM not loaded or library unavailable." if image_pil is None: return "Please upload an image." print("Running FastSAM 'segment everything'...") try: if image_pil.mode != "RGB": image_pil = image_pil.convert("RGB") image_np_rgb = np.array(image_pil) everything_results = fastsam_model( image_np_rgb, device=DEVICE, retina_masks=True, imgsz=640, conf=conf_threshold, iou=iou_threshold, verbose=True ) prompt_process = FastSAMPrompt(image_np_rgb, everything_results, device=DEVICE) ann = prompt_process.everything_prompt() output_image = image_pil.copy() if ann and ann[0] and 'masks' in ann[0] and len(ann[0]['masks']) > 0: masks = ann[0]['masks'].cpu().numpy() print(f"Found {len(masks)} masks with shape: {masks.shape}") overlay = Image.new('RGBA', output_image.size, (0, 0, 0, 0)) draw = ImageDraw.Draw(overlay) for mask in masks: mask = (mask > 0).astype(np.uint8) * 255 color = (random.randint(50, 255), random.randint(50, 255), random.randint(50, 255), 180) mask_image = Image.fromarray(mask, mode='L') draw.bitmap((0, 0), mask_image, fill=color) output_image = Image.alpha_composite(output_image.convert('RGBA'), overlay).convert('RGB') else: print("No masks detected in 'segment everything' mode.") return output_image except Exception as e: print(f"Error during FastSAM 'everything' processing: {e}") traceback.print_exc() return f"Error: {e}" def run_text_prompted_segmentation(image_pil: Image.Image, text_prompts: str, conf_threshold: float = 0.4, iou_threshold: float = 0.9): if not load_fastsam_model(): return "Error: FastSAM Model not loaded.", "Model load failure." if not fastsam_lib_imported: return "Error: FastSAM library not available.", "Library import error." if image_pil is None: return "Please upload an image.", "No image provided." if not text_prompts: return image_pil, "Please enter text prompts (e.g., 'person, dog')." prompts = [p.strip() for p in text_prompts.split(',') if p.strip()] if not prompts: return image_pil, "No valid text prompts entered." print(f"Running FastSAM text-prompted segmentation for: {prompts}") try: if image_pil.mode != "RGB": image_pil = image_pil.convert("RGB") image_np_rgb = np.array(image_pil) everything_results = fastsam_model( image_np_rgb, device=DEVICE, retina_masks=True, imgsz=640, conf=conf_threshold, iou=iou_threshold, verbose=True ) prompt_process = FastSAMPrompt(image_np_rgb, everything_results, device=DEVICE) all_matching_masks = [] found_prompts = [] for text in prompts: print(f" Processing prompt: '{text}'") ann = prompt_process.text_prompt(text=text) if ann and ann[0] and 'masks' in ann[0] and len(ann[0]['masks']) > 0: num_found = len(ann[0]['masks']) print(f" Found {num_found} mask(s) with shape: {ann[0]['masks'].shape}") found_prompts.append(f"{text} ({num_found})") masks = ann[0]['masks'].cpu().numpy() all_matching_masks.extend(masks) else: print(f" No masks found for '{text}'.") found_prompts.append(f"{text} (0)") output_image = image_pil.copy() status_message = f"Found segments for: {', '.join(found_prompts)}" if found_prompts else "No matches found." if all_matching_masks: masks_np = np.stack(all_matching_masks, axis=0) print(f"Total masks stacked: {masks_np.shape}") overlay = Image.new('RGBA', output_image.size, (0, 0, 0, 0)) draw = ImageDraw.Draw(overlay) for mask in masks_np: mask = (mask > 0).astype(np.uint8) * 255 color = (random.randint(50, 255), random.randint(50, 255), random.randint(50, 255), 180) mask_image = Image.fromarray(mask, mode='L') draw.bitmap((0, 0), mask_image, fill=color) output_image = Image.alpha_composite(output_image.convert('RGBA'), overlay).convert('RGB') return output_image, status_message except Exception as e: print(f"Error during FastSAM text-prompted processing: {e}") traceback.print_exc() return image_pil, f"Error: {e}" # --- Gradio Interface --- print("Attempting to preload models...") load_fastsam_model() # Load FastSAM eagerly print("Preloading finished.") with gr.Blocks(theme=gr.themes.Soft()) as demo: gr.Markdown("# CLIP & FastSAM Demo") gr.Markdown("Explore Zero-Shot Classification, 'Segment Everything', and Text-Prompted Segmentation.") with gr.Tabs(): with gr.TabItem("CLIP Zero-Shot Classification"): gr.Markdown("Upload an image and provide comma-separated labels (e.g., 'cat, dog, car').") with gr.Row(): with gr.Column(scale=1): clip_input_image = gr.Image(type="pil", label="Input Image") clip_text_labels = gr.Textbox(label="Comma-Separated Labels", placeholder="e.g., astronaut, moon") clip_button = gr.Button("Run CLIP Classification", variant="primary") with gr.Column(scale=1): clip_output_label = gr.Label(label="Classification Probabilities") clip_output_image_display = gr.Image(type="pil", label="Input Image Preview") clip_button.click( run_clip_zero_shot, inputs=[clip_input_image, clip_text_labels], outputs=[clip_output_label, clip_output_image_display] ) gr.Examples( examples=[ ["examples/astronaut.jpg", "astronaut, moon, rover"], ["examples/dog_bike.jpg", "dog, bicycle, person"], ["examples/clip_logo.png", "logo, text, graphics"], ], inputs=[clip_input_image, clip_text_labels], outputs=[clip_output_label, clip_output_image_display], fn=run_clip_zero_shot, cache_examples=False, ) with gr.TabItem("FastSAM Segment Everything"): gr.Markdown("Upload an image to segment all objects/regions.") with gr.Row(): with gr.Column(scale=1): fastsam_input_image_all = gr.Image(type="pil", label="Input Image") with gr.Row(): fastsam_conf_all = gr.Slider(minimum=0.1, maximum=1.0, value=0.4, step=0.05, label="Confidence Threshold") fastsam_iou_all = gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="IoU Threshold") fastsam_button_all = gr.Button("Run FastSAM Segmentation", variant="primary") with gr.Column(scale=1): fastsam_output_image_all = gr.Image(type="pil", label="Segmented Image") fastsam_button_all.click( run_fastsam_segmentation, inputs=[fastsam_input_image_all, fastsam_conf_all, fastsam_iou_all], outputs=[fastsam_output_image_all] ) gr.Examples( examples=[ ["examples/dogs.jpg", 0.4, 0.9], ["examples/fruits.jpg", 0.5, 0.8], ["examples/lion.jpg", 0.45, 0.9], ], inputs=[fastsam_input_image_all, fastsam_conf_all, fastsam_iou_all], outputs=[fastsam_output_image_all], fn=run_fastsam_segmentation, cache_examples=False, ) with gr.TabItem("Text-Prompted Segmentation"): gr.Markdown("Upload an image and provide comma-separated prompts (e.g., 'person, dog').") with gr.Row(): with gr.Column(scale=1): prompt_input_image = gr.Image(type="pil", label="Input Image") prompt_text_input = gr.Textbox(label="Comma-Separated Text Prompts", placeholder="e.g., glasses, watch") with gr.Row(): prompt_conf = gr.Slider(minimum=0.1, maximum=1.0, value=0.4, step=0.05, label="Confidence Threshold") prompt_iou = gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="IoU Threshold") prompt_button = gr.Button("Segment by Text", variant="primary") with gr.Column(scale=1): prompt_output_image = gr.Image(type="pil", label="Text-Prompted Segmentation") prompt_status_message = gr.Textbox(label="Status", interactive=False) prompt_button.click( run_text_prompted_segmentation, inputs=[prompt_input_image, prompt_text_input, prompt_conf, prompt_iou], outputs=[prompt_output_image, prompt_status_message] ) gr.Examples( examples=[ ["examples/dog_bike.jpg", "person, bicycle", 0.4, 0.9], ["examples/astronaut.jpg", "person, helmet", 0.35, 0.9], ["examples/dogs.jpg", "dog", 0.4, 0.9], ["examples/fruits.jpg", "banana, apple", 0.5, 0.8], ["examples/teacher.jpg", "person, glasses", 0.4, 0.9], ], inputs=[prompt_input_image, prompt_text_input, prompt_conf, prompt_iou], outputs=[prompt_output_image, prompt_status_message], fn=run_text_prompted_segmentation, cache_examples=False, ) # Download example images with retries if not os.path.exists("examples"): os.makedirs("examples") print("Created 'examples' directory.") example_files = { "astronaut.jpg": "https://upload.wikimedia.org/wikipedia/commons/thumb/d/d1/Astronaut_-_St._Jean_Bay.jpg/640px-Astronaut_-_St._Jean_Bay.jpg", "dog_bike.jpg": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/gradio/outputs_multimodal.jpg", "clip_logo.png": "https://raw.githubusercontent.com/openai/CLIP/main/CLIP.png", "dogs.jpg": "https://raw.githubusercontent.com/ultralytics/assets/main/im/image8.jpg", "fruits.jpg": "https://raw.githubusercontent.com/ultralytics/assets/main/im/image9.jpg", "lion.jpg": "https://huggingface.co/spaces/gradio/image-segmentation/resolve/main/images/lion.jpg", "teacher.jpg": "https://images.pexels.com/photos/848117/pexels-photo-848117.jpeg?auto=compress&cs=tinysrgb&w=600" } def download_example_file(filename, url, retries=3): filepath = os.path.join("examples", filename) if not os.path.exists(filepath): for attempt in range(retries): try: print(f"Downloading {filename} (attempt {attempt + 1}/{retries})...") wget.download(url, filepath) break except Exception as e: print(f"Attempt {attempt + 1} failed: {e}") if attempt + 1 == retries: print(f"Failed to download {filename} after {retries} attempts.") for filename, url in example_files.items(): download_example_file(filename, url) if __name__ == "__main__": demo.launch(debug=True)