app.py

import gradio as gr
import torch
import cv2
import numpy as np
from transformers import SamModel, SamProcessor, BlipProcessor, BlipForConditionalGeneration
from PIL import Image

# Set up device
device = "cuda" if torch.cuda.is_available() else "cpu"

# Load SAM model and processor
sam_model = SamModel.from_pretrained("facebook/sam-vit-base").to(device)
sam_processor = SamProcessor.from_pretrained("facebook/sam-vit-base")

# Load BLIP model and processor for image-to-text
blip_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
blip_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base").to(device)

def process_mask(mask, target_size):
    if mask.ndim > 2:
        mask = mask.squeeze()
    if mask.ndim > 2:
        mask = mask[0]
    mask = (mask > 0.5).astype(np.uint8) * 255
    mask_image = Image.fromarray(mask)
    mask_image = mask_image.resize(target_size, Image.NEAREST)
    return np.array(mask_image) > 0

def segment_image(input_image, object_name):
    try:
        if input_image is None:
            return None, "Please upload an image before submitting."
        
        input_image = Image.fromarray(input_image).convert("RGB")
        original_size = input_image.size
        if not original_size or 0 in original_size:
            return None, "Invalid image size. Please upload a different image."
        
        # Generate image caption
        blip_inputs = blip_processor(input_image, return_tensors="pt").to(device)
        caption = blip_model.generate(**blip_inputs)
        caption_text = blip_processor.decode(caption[0], skip_special_tokens=True)
        
        # Process the image with SAM
        sam_inputs = sam_processor(input_image, return_tensors="pt").to(device)
        
        # Generate masks
        with torch.no_grad():
            sam_outputs = sam_model(**sam_inputs)
        
        # Post-process masks
        masks = sam_processor.image_processor.post_process_masks(
            sam_outputs.pred_masks.cpu(),
            sam_inputs["original_sizes"].cpu(),
            sam_inputs["reshaped_input_sizes"].cpu()
        )
        
        # Find the mask that best matches the specified object
        best_mask = None
        best_score = -1
        for mask in masks[0]:
            mask_binary = mask.numpy() > 0.5
            mask_area = mask_binary.sum()
            if object_name.lower() in caption_text.lower() and mask_area > best_score:
                best_mask = mask_binary
                best_score = mask_area
        
        if best_mask is None:
            return input_image, f"Could not find '{object_name}' in the image."
        
        combined_mask = process_mask(best_mask, original_size)
        
        # Overlay the mask on the original image
        result_image = np.array(input_image)
        mask_rgb = np.zeros_like(result_image)
        mask_rgb[combined_mask] = [255, 0, 0]  # Red color for the mask
        result_image = cv2.addWeighted(result_image, 1, mask_rgb, 0.5, 0)
        
        return result_image, f"Segmented '{object_name}' in the image."
    
    except Exception as e:
        return None, f"An error occurred: {str(e)}"

# Create Gradio interface
iface = gr.Interface(
    fn=segment_image,
    inputs=[
        gr.Image(type="numpy", label="Upload an image"),
        gr.Textbox(label="Specify object to segment (e.g., dog, cat, grass)")
    ],
    outputs=[
        gr.Image(type="numpy", label="Segmented Image"),
        gr.Textbox(label="Status")
    ],
    title="Segment Anything Model (SAM) with Object Specification",
    description="Upload an image and specify an object to segment."
)

# Launch the interface
iface.launch()