app.py

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)