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 # To download weights

# --- 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 ---
# Use a smaller model suitable for Spaces CPU/basic GPU if needed
FASTSAM_CHECKPOINT = "FastSAM-s.pt"
FASTSAM_CHECKPOINT_URL = f"https://huggingface.co/spaces/An-619/FastSAM/resolve/main/{FASTSAM_CHECKPOINT}" # Example URL, find official if possible

fastsam_model = None

def download_fastsam_weights():
    if not os.path.exists(FASTSAM_CHECKPOINT):
        print(f"Downloading FastSAM weights: {FASTSAM_CHECKPOINT}...")
        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.")
            return False
    return os.path.exists(FASTSAM_CHECKPOINT)

def load_fastsam_model():
    global fastsam_model
    if fastsam_model is None:
        if download_fastsam_weights():
            try:
                from fastsam import FastSAM, FastSAMPrompt # Import here after potential download
                print(f"Loading FastSAM model: {FASTSAM_CHECKPOINT}...")
                fastsam_model = FastSAM(FASTSAM_CHECKPOINT)
                print(f"FastSAM model loaded.") # Device handled internally by FastSAM based on its setup/torch device
            except ImportError:
                print("Error: 'fastsam' library not found. Please install it (pip install fastsam).")
            except Exception as e:
                print(f"Error loading FastSAM model: {e}")
        else:
            print("FastSAM weights not found. 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 not text_labels:
        return "Please provide comma-separated text labels.", None
    if image is None:
        return "Please upload an image.", None

    labels = [label.strip() for label in text_labels.split(',')]
    if not labels:
         return "No valid labels provided.", None

    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  # this is the image-text similarity score
            probs = logits_per_image.softmax(dim=1)  # convert to probabilities

        print("CLIP processing complete.")

        # Format output for Gradio Label
        confidences = {labels[i]: float(probs[0, i].item()) for i in range(len(labels))}
        return confidences, image # Return original image for display alongside results

    except Exception as e:
        print(f"Error during CLIP processing: {e}")
        return f"An error occurred: {e}", None


# FastSAM Segmentation Function
def run_fastsam_segmentation(image_pil: Image.Image, conf_threshold: float = 0.4, iou_threshold: float = 0.9):
    if fastsam_model is None:
        load_fastsam_model() # Attempt to load if not already loaded
        if fastsam_model is None:
             return "Error: FastSAM Model not loaded. Check logs.", None
    if image_pil is None:
        return "Please upload an image.", None

    print("Running FastSAM segmentation...")

    try:
         # Ensure image is RGB
        if image_pil.mode != "RGB":
            image_pil = image_pil.convert("RGB")

        # FastSAM expects a BGR numpy array or path usually. Let's try with RGB numpy.
        # If it fails, uncomment the BGR conversion line.
        image_np_rgb = np.array(image_pil)
        # image_np_bgr = image_np_rgb[:, :, ::-1] # Convert RGB to BGR if needed

        # Run FastSAM inference
        # Adjust imgsz, conf, iou as needed. Higher imgsz = more detail, slower.
        everything_results = fastsam_model(
            image_np_rgb, # Use image_np_bgr if conversion needed
            device=DEVICE,
            retina_masks=True,
            imgsz=640, # Smaller size for faster inference on limited hardware
            conf=conf_threshold,
            iou=iou_threshold,
        )

        # Process results using FastSAMPrompt
        from fastsam import FastSAMPrompt # Make sure it's imported
        prompt_process = FastSAMPrompt(image_np_rgb, everything_results, device=DEVICE)

        # Get all annotations (masks)
        ann = prompt_process.everything_prompt()

        print(f"FastSAM found {len(ann[0]['masks']) if ann and ann[0] else 0} masks.")

        # --- Plotting Masks on Image (Manual) ---
        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() # shape (N, H, W)
            
            # Create overlay image
            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] # shape (H, W), boolean
                
                # Generate random color with some transparency
                color = (random.randint(50, 255), random.randint(50, 255), random.randint(50, 255), 128) # RGBA with alpha

                # Create a single-channel image from the boolean mask
                mask_image = Image.fromarray((mask * 255).astype(np.uint8), mode='L')
                
                # Apply color to the mask area on the overlay
                draw.bitmap((0,0), mask_image, fill=color)

            # Composite the overlay onto the original image
            output_image = Image.alpha_composite(output_image.convert('RGBA'), overlay).convert('RGB')

        print("FastSAM processing and plotting complete.")
        return output_image, image_pil # Return segmented and original images

    except Exception as e:
        print(f"Error during FastSAM processing: {e}")
        import traceback
        traceback.print_exc() # Print detailed traceback
        return f"An error occurred: {e}", None


# --- Gradio Interface ---

# Pre-load models on startup (optional but good for performance)
print("Attempting to preload models...")
load_clip_model()
load_fastsam_model()
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, mountain, 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") # Show input for context

            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"],
                ],
                inputs=[clip_input_image, clip_text_labels],
                outputs=[clip_output_label, clip_output_image_display],
                fn=run_clip_zero_shot,
                cache_examples=False, # Re-run for live demo
            )

        # --- 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_input_display = gr.Image(type="pil", label="Original Image") # Optional: show original side-by-side

            fastsam_button.click(
                run_fastsam_segmentation,
                inputs=[fastsam_input_image, fastsam_conf, fastsam_iou],
                outputs=[fastsam_output_image] # Removed the second output for simplicity, adjust if needed
            )
            gr.Examples(
                examples=[
                    ["examples/dogs.jpg", 0.4, 0.9],
                    ["examples/fruits.jpg", 0.5, 0.8],
                ],
                inputs=[fastsam_input_image, fastsam_conf, fastsam_iou],
                outputs=[fastsam_output_image],
                fn=run_fastsam_segmentation,
                cache_examples=False, # Re-run for live demo
            )

    # Add example images (optional, but helpful)
    # Create an 'examples' folder and add some jpg images like 'astronaut.jpg', 'dog_bike.jpg', 'dogs.jpg', 'fruits.jpg'
    if not os.path.exists("examples"):
        os.makedirs("examples")
        print("Created 'examples' directory. Please add some images (e.g., astronaut.jpg, dog_bike.jpg) for the examples to work.")
        # You might need to download some sample images here too if running on a fresh env
        try:
            print("Downloading example images...")
            wget.download("https://huggingface.co/spaces/gradio/image-segmentation/resolve/main/images/lion.jpg", "examples/lion.jpg")
            wget.download("https://raw.githubusercontent.com/openai/CLIP/main/CLIP.png", "examples/clip_logo.png")
            wget.download("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/gradio-logo.png", "examples/gradio_logo.png")
            # Manually add the examples used above if these don't match
            print("Example images downloaded (or attempted). Please verify.")
        except Exception as e:
             print(f"Could not download example images: {e}")


# Launch the Gradio app
if __name__ == "__main__":
    demo.launch(debug=True) # Set debug=False for deployment