app.py

import gradio as gr
import torch
from PIL import Image
import cv2
import numpy as np
from transformers import CLIPProcessor, CLIPModel
from ultralytics import FastSAM
import supervision as sv
import os

# Load CLIP model
model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")

# Initialize FastSAM model
FASTSAM_WEIGHTS = "FastSAM-s.pt"
if not os.path.exists(FASTSAM_WEIGHTS):
    os.system(f"wget https://huggingface.co/spaces/An-619/FastSAM/resolve/main/weights/{FASTSAM_WEIGHTS}")
fast_sam = FastSAM(FASTSAM_WEIGHTS)

def process_image_clip(image, text_input):
    if image is None:
        return "Please upload an image first."
    if not text_input:
        return "Please enter some text to check in the image."
    
    try:
        # Convert numpy array to PIL Image if needed
        if isinstance(image, np.ndarray):
            image = Image.fromarray(image)
        
        # Create a list of candidate labels
        candidate_labels = [text_input, f"not {text_input}"]
        
        # Process image and text
        inputs = processor(
            images=image,
            text=candidate_labels,
            return_tensors="pt",
            padding=True
        )
        
        # Get model predictions
        outputs = model(**{k: v for k, v in inputs.items()})
        logits_per_image = outputs.logits_per_image
        probs = logits_per_image.softmax(dim=1)
        
        # Get confidence for the positive label
        confidence = float(probs[0][0])
        return f"Confidence that the image contains '{text_input}': {confidence:.2%}"
    except Exception as e:
        return f"Error processing image: {str(e)}"

def process_image_fastsam(image):
    if image is None:
        return None
        
    try:
        # Convert PIL image to numpy array if needed
        if isinstance(image, Image.Image):
            image_np = np.array(image)
        else:
            image_np = image
            
        # Run FastSAM inference
        results = fast_sam(image_np, device='cpu', retina_masks=True, imgsz=1024, conf=0.4, iou=0.9)
        
        # Get detections
        detections = sv.Detections.from_ultralytics(results[0])
        
        # Create annotator
        box_annotator = sv.BoxAnnotator()
        mask_annotator = sv.MaskAnnotator()
        
        # Annotate image
        annotated_image = mask_annotator.annotate(scene=image_np.copy(), detections=detections)
        annotated_image = box_annotator.annotate(scene=annotated_image, detections=detections)
        
        return Image.fromarray(annotated_image)
    except Exception as e:
        return f"Error processing image: {str(e)}"

# Create Gradio interface
with gr.Blocks(css="footer {visibility: hidden}") as demo:
    gr.Markdown("""
    # CLIP and FastSAM Demo
    This demo combines two powerful AI models:
    - **CLIP**: For zero-shot image classification
    - **FastSAM**: For automatic image segmentation
    
    Try uploading an image and use either of the tabs below!
    """)
    
    with gr.Tab("CLIP Zero-Shot Classification"):
        with gr.Row():
            image_input = gr.Image(label="Input Image")
            text_input = gr.Textbox(
                label="What do you want to check in the image?", 
                placeholder="e.g., 'a dog', 'sunset', 'people playing'",
                info="Enter any concept you want to check in the image"
            )
        output_text = gr.Textbox(label="Result")
        classify_btn = gr.Button("Classify")
        classify_btn.click(fn=process_image_clip, inputs=[image_input, text_input], outputs=output_text)
        
        gr.Examples(
            examples=[
                ["https://raw.githubusercontent.com/gradio-app/gradio/main/demo/kitchen/kitchen.png", "kitchen"],
                ["https://raw.githubusercontent.com/gradio-app/gradio/main/demo/calculator/calculator.jpg", "calculator"],
            ],
            inputs=[image_input, text_input],
        )
    
    with gr.Tab("FastSAM Segmentation"):
        with gr.Row():
            image_input_sam = gr.Image(label="Input Image")
            image_output = gr.Image(label="Segmentation Result")
        segment_btn = gr.Button("Segment")
        segment_btn.click(fn=process_image_fastsam, inputs=[image_input_sam], outputs=image_output)
        
        gr.Examples(
            examples=[
                ["https://raw.githubusercontent.com/gradio-app/gradio/main/demo/kitchen/kitchen.png"],
                ["https://raw.githubusercontent.com/gradio-app/gradio/main/demo/calculator/calculator.jpg"],
            ],
            inputs=[image_input_sam],
        )
    
    gr.Markdown("""
    ### How to use:
    1. **CLIP Classification**: Upload an image and enter text to check if that concept exists in the image
    2. **FastSAM Segmentation**: Upload an image to get automatic segmentation with bounding boxes and masks
    
    ### Note:
    - The models run on CPU, so processing might take a few seconds
    - For best results, use clear images with good lighting
    """)

demo.launch(share=True)