from PIL import Image
from transformers import CLIPProcessor, CLIPModel
import gradio as gr
import torchvision.transforms as transforms

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

def image_similarity(image: Image.Image, positive_prompt: str, negative_prompts: str):
    # Convert the PIL Image to a tensor and preprocess
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
    ])
    image_tensor = transform(image).unsqueeze(0)  # Add batch dimension

    # Split the negative prompts string into a list of prompts
    negative_prompts_list = negative_prompts.split(";")
    # Combine positive and negative prompts into one list
    prompts = [positive_prompt.strip()] + [np.strip() for np in negative_prompts_list]

    # Process prompts and image tensor
    inputs = processor(
        text=prompts,
        images=image_tensor,
        return_tensors="pt",
        padding=True
    )

    outputs = model(**inputs)
    logits_per_image = outputs.logits_per_image
    probs = logits_per_image.softmax(dim=1)

    # Determine if positive prompt has a higher probability than any of the negative prompts
    is_positive_highest = probs[0][0] > max(probs[0][1:])

    return bool(is_positive_highest), f"Probability for Positive Prompt: {probs[0][0]:.4f}"

interface = gr.Interface(
    fn=image_similarity, 
    inputs=[
        gr.components.Image(type="pil"), 
        gr.components.Text(label="Enter positive prompt e.g. 'a person drinking a beverage'"),
        gr.components.Textbox(label="Enter negative prompts, separated by semicolon e.g. 'an empty scene; person without beverage'", placeholder="negative prompt 1; negative prompt 2; ..."),
    ], 
    outputs=[
        gr.components.Textbox(label="Result"),
        gr.components.Textbox(label="Probability for Positive Prompt")
    ],
    title="Engagify's Image Action Detection",
    description="[Author: Ibrahim Hasani] This Method uses CLIP-VIT [Version: BASE-PATCH-16] to determine if an action is being performed in an image or not. (Binary Classifier). It contrasts an Action against multiple negative labels. Ensure the prompts accurately describe the desired