import gradio as gr
from transformers import CLIPProcessor, CLIPModel
from PIL import Image
import torch

MODEL_ID = "KhanyiTapiwa00/medsiglip-diagnosis"

processor = CLIPProcessor.from_pretrained(MODEL_ID)
model = CLIPModel.from_pretrained(MODEL_ID)
model.eval()

def predict(image: Image.Image, text: str):
    if image is None or text.strip() == "":
        return "Please provide both an image and a text description."
    
    inputs = processor(images=image, text=text, return_tensors="pt", padding=True)
    with torch.no_grad():
        outputs = model(**inputs)
    
    logits = outputs.logits_per_image
    probs = torch.softmax(logits, dim=1)
    return str(probs.cpu().numpy())

demo = gr.Interface(
    fn=predict,
    inputs=[gr.Image(type="pil", label="Upload Medical Image"),
            gr.Textbox(label="Enter Description")],
    outputs=gr.Textbox(label="Similarity Score"),
    title="MedSigLIP AI Demo",
    description="Upload a medical image and compare it with a text description."
)

demo.launch()