Spaces:
Sleeping
Sleeping
| 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 # To download weights | |
| import traceback # For detailed error printing | |
| # --- Configuration & Model Loading --- | |
| # Device Selection | |
| DEVICE = "cuda" if torch.cuda.is_available() 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: | |
| print(f"Loading CLIP processor: {CLIP_MODEL_ID}...") | |
| clip_processor = AutoProcessor.from_pretrained(CLIP_MODEL_ID) | |
| print("CLIP processor loaded.") | |
| if clip_model is None: | |
| 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}.") | |
| # --- FastSAM Setup --- | |
| FASTSAM_CHECKPOINT = "FastSAM-s.pt" | |
| # Use the official model hub repo URL | |
| FASTSAM_CHECKPOINT_URL = f"https://huggingface.co/CASIA-IVA-Lab/FastSAM-s/resolve/main/{FASTSAM_CHECKPOINT}" | |
| fastsam_model = None | |
| fastsam_lib_imported = False # Flag to check if import worked | |
| def check_and_import_fastsam(): | |
| global fastsam_lib_imported | |
| if not fastsam_lib_imported: | |
| try: | |
| from fastsam import FastSAM, FastSAMPrompt | |
| globals()['FastSAM'] = FastSAM # Make classes available globally | |
| globals()['FastSAMPrompt'] = FastSAMPrompt | |
| fastsam_lib_imported = True | |
| print("fastsam library imported successfully.") | |
| except ImportError: | |
| print("Error: 'fastsam' library not found or import failed.") | |
| print("Please ensure 'fastsam' is installed correctly (pip install fastsam).") | |
| fastsam_lib_imported = False | |
| except Exception as e: | |
| print(f"An unexpected error occurred during fastsam import: {e}") | |
| fastsam_lib_imported = False | |
| return fastsam_lib_imported | |
| def download_fastsam_weights(): | |
| if not os.path.exists(FASTSAM_CHECKPOINT): | |
| print(f"Downloading FastSAM weights: {FASTSAM_CHECKPOINT} from {FASTSAM_CHECKPOINT_URL}...") | |
| try: | |
| wget.download(FASTSAM_CHECKPOINT_URL, FASTSAM_CHECKPOINT) | |
| print("FastSAM weights downloaded.") | |
| except Exception as e: | |
| print(f"Error downloading FastSAM weights: {e}") | |
| print("Please ensure the URL is correct and reachable, or manually place the weights file.") | |
| # Attempt to remove partially downloaded file if exists | |
| if os.path.exists(FASTSAM_CHECKPOINT): | |
| try: | |
| os.remove(FASTSAM_CHECKPOINT) | |
| except OSError: | |
| pass # Ignore removal errors | |
| 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(): # Check import first | |
| print("Cannot load FastSAM model because the library couldn't be imported.") | |
| return # Exit if import failed | |
| if download_fastsam_weights(): # Check download/existence second | |
| try: | |
| # FastSAM class should be available via globals() now | |
| print(f"Loading FastSAM model: {FASTSAM_CHECKPOINT}...") | |
| fastsam_model = FastSAM(FASTSAM_CHECKPOINT) | |
| print(f"FastSAM model loaded.") # Device handled internally by FastSAM | |
| except Exception as e: | |
| print(f"Error loading FastSAM model: {e}") | |
| traceback.print_exc() | |
| else: | |
| print("FastSAM weights not found or download failed. Cannot load model.") | |
| # --- Processing Functions --- | |
| # CLIP Zero-Shot Classification Function | |
| def run_clip_zero_shot(image: Image.Image, text_labels: str): | |
| if clip_model is None or clip_processor is None: | |
| load_clip_model() # Attempt to load if not already loaded | |
| if clip_model is None: | |
| return "Error: CLIP Model not loaded. Check logs.", None | |
| if image is None: | |
| return "Please upload an image.", None # Return None for the image display | |
| if not text_labels: | |
| # Return empty results but display the uploaded image | |
| return {}, image | |
| labels = [label.strip() for label in text_labels.split(',') if label.strip()] # Ensure non-empty labels | |
| if not labels: | |
| # Return empty results but display the uploaded image | |
| return {}, image | |
| print(f"Running CLIP zero-shot classification with labels: {labels}") | |
| try: | |
| # Ensure image is RGB | |
| 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) | |
| logits_per_image = outputs.logits_per_image | |
| probs = logits_per_image.softmax(dim=1) | |
| print("CLIP processing complete.") | |
| confidences = {labels[i]: float(probs[0, i].item()) for i in range(len(labels))} | |
| # Return results and the original image used for prediction | |
| return confidences, image | |
| except Exception as e: | |
| print(f"Error during CLIP processing: {e}") | |
| traceback.print_exc() | |
| # Return error message and the original image | |
| return f"An error occurred during CLIP: {e}", image | |
| # FastSAM Segmentation Function | |
| def run_fastsam_segmentation(image_pil: Image.Image, conf_threshold: float = 0.4, iou_threshold: float = 0.9): | |
| # Ensure model is loaded or attempt to load | |
| if fastsam_model is None: | |
| load_fastsam_model() | |
| if fastsam_model is None: | |
| # Return error message string for the image component (Gradio handles this) | |
| return "Error: FastSAM Model not loaded. Check logs." | |
| # Ensure library was imported | |
| if not fastsam_lib_imported: | |
| return "Error: FastSAM library not available. Cannot run segmentation." | |
| if image_pil is None: | |
| return "Please upload an image." | |
| print("Running FastSAM segmentation...") | |
| try: | |
| # Ensure image is RGB | |
| if image_pil.mode != "RGB": | |
| image_pil = image_pil.convert("RGB") | |
| image_np_rgb = np.array(image_pil) | |
| # Run FastSAM inference | |
| everything_results = fastsam_model( | |
| image_np_rgb, | |
| device=DEVICE, | |
| retina_masks=True, | |
| imgsz=640, | |
| conf=conf_threshold, | |
| iou=iou_threshold, | |
| ) | |
| # FastSAMPrompt should be available via globals() if import succeeded | |
| prompt_process = FastSAMPrompt(image_np_rgb, everything_results, device=DEVICE) | |
| ann = prompt_process.everything_prompt() | |
| print(f"FastSAM found {len(ann[0]['masks']) if ann and ann[0] and 'masks' in ann[0] else 0} masks.") | |
| # --- Plotting Masks on Image --- | |
| output_image = image_pil.copy() | |
| if ann and ann[0] is not None and 'masks' in ann[0] and len(ann[0]['masks']) > 0: | |
| masks = ann[0]['masks'].cpu().numpy() # (N, H, W) boolean | |
| overlay = Image.new('RGBA', output_image.size, (0, 0, 0, 0)) | |
| draw = ImageDraw.Draw(overlay) | |
| for i in range(masks.shape[0]): | |
| mask = masks[i] | |
| color = (random.randint(50, 255), random.randint(50, 255), random.randint(50, 255), 128) # RGBA | |
| mask_image = Image.fromarray((mask * 255).astype(np.uint8), mode='L') | |
| draw.bitmap((0,0), mask_image, fill=color) | |
| output_image = Image.alpha_composite(output_image.convert('RGBA'), overlay).convert('RGB') | |
| print("FastSAM processing and plotting complete.") | |
| # *** FIX: Return ONLY the output image for the single Image component *** | |
| return output_image | |
| except NameError as ne: | |
| print(f"NameError during FastSAM processing: {ne}. Was the fastsam library imported correctly?") | |
| traceback.print_exc() | |
| return f"A NameError occurred: {ne}. Check library import." | |
| except Exception as e: | |
| print(f"Error during FastSAM processing: {e}") | |
| traceback.print_exc() | |
| return f"An error occurred during FastSAM: {e}" | |
| # --- Gradio Interface --- | |
| # Pre-load models on startup (optional but good for performance) | |
| print("Attempting to preload models...") | |
| load_clip_model() | |
| load_fastsam_model() # This will now also attempt download/check import | |
| print("Preloading finished (or attempted).") | |
| with gr.Blocks(theme=gr.themes.Soft()) as demo: | |
| gr.Markdown("# CLIP & FastSAM Demo") | |
| gr.Markdown("Explore Zero-Shot Classification with CLIP and 'Segment Anything' with FastSAM.") | |
| with gr.Tabs(): | |
| # --- CLIP Tab --- | |
| with gr.TabItem("CLIP Zero-Shot Classification"): | |
| gr.Markdown("Upload an image and provide comma-separated candidate labels (e.g., 'cat, dog, car'). CLIP will predict the probability of the image matching each label.") | |
| 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, dog playing fetch") | |
| 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, mountain"], | |
| ["examples/dog_bike.jpg", "dog, bicycle, person, park, grass"], | |
| ["examples/clip_logo.png", "logo, text, graphics, abstract art"], # Added another example | |
| ], | |
| inputs=[clip_input_image, clip_text_labels], | |
| outputs=[clip_output_label, clip_output_image_display], | |
| fn=run_clip_zero_shot, | |
| cache_examples=False, | |
| ) | |
| # --- FastSAM Tab --- | |
| with gr.TabItem("FastSAM Segmentation"): | |
| gr.Markdown("Upload an image. FastSAM will attempt to segment all objects/regions in the image.") | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| fastsam_input_image = gr.Image(type="pil", label="Input Image") | |
| with gr.Row(): | |
| fastsam_conf = gr.Slider(minimum=0.1, maximum=1.0, value=0.4, step=0.05, label="Confidence Threshold") | |
| fastsam_iou = gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="IoU Threshold") | |
| fastsam_button = gr.Button("Run FastSAM Segmentation", variant="primary") | |
| with gr.Column(scale=1): | |
| fastsam_output_image = gr.Image(type="pil", label="Segmented Image") | |
| fastsam_button.click( | |
| run_fastsam_segmentation, | |
| inputs=[fastsam_input_image, fastsam_conf, fastsam_iou], | |
| # Output is now correctly mapped to the single component | |
| outputs=[fastsam_output_image] | |
| ) | |
| gr.Examples( | |
| examples=[ | |
| ["examples/dogs.jpg", 0.4, 0.9], | |
| ["examples/fruits.jpg", 0.5, 0.8], | |
| ["examples/lion.jpg", 0.45, 0.9], # Added another example | |
| ], | |
| inputs=[fastsam_input_image, fastsam_conf, fastsam_iou], | |
| outputs=[fastsam_output_image], | |
| fn=run_fastsam_segmentation, | |
| cache_examples=False, | |
| ) | |
| # Add example images (optional, but helpful) | |
| if not os.path.exists("examples"): | |
| os.makedirs("examples") | |
| print("Created 'examples' directory. Attempting to download sample images...") | |
| example_files = { | |
| "astronaut.jpg": "https://upload.wikimedia.org/wikipedia/commons/thumb/d/d astronaut_-_St._Jean_Bay.jpg/640px-Astronaut_-_St._Jean_Bay.jpg", # Find suitable public domain/CC image | |
| "dog_bike.jpg": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/gradio/outputs_multimodal.jpg", # Using a relevant example from HF | |
| "clip_logo.png": "https://raw.githubusercontent.com/openai/CLIP/main/CLIP.png", | |
| "dogs.jpg": "https://raw.githubusercontent.com/ultralytics/assets/main/im/image8.jpg", # From Ultralytics assets | |
| "fruits.jpg": "https://raw.githubusercontent.com/ultralytics/assets/main/im/image9.jpg", # From Ultralytics assets | |
| "lion.jpg": "https://huggingface.co/spaces/gradio/image-segmentation/resolve/main/images/lion.jpg" | |
| } | |
| for filename, url in example_files.items(): | |
| filepath = os.path.join("examples", filename) | |
| if not os.path.exists(filepath): | |
| try: | |
| print(f"Downloading {filename}...") | |
| wget.download(url, filepath) | |
| except Exception as e: | |
| print(f"Could not download {filename} from {url}: {e}") | |
| print("Example image download attempt finished.") | |
| # Launch the Gradio app | |
| if __name__ == "__main__": | |
| # share=True is primarily for local testing to get a public link. | |
| # Not needed/used when deploying on Hugging Face Spaces. | |
| # debug=True is helpful for development. Set to False for production. | |
| demo.launch(debug=True) |