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

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

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

# Load normal image for reference
normal_image = Image.open("normal_sample.jpg")

with torch.no_grad():
    inputs = processor(images=normal_image, return_tensors="pt").to(device)
    normal_features = model.get_image_features(**inputs)
    normal_features = normal_features / normal_features.norm(p=2, dim=-1, keepdim=True)

def detect_anomaly(img):
    with torch.no_grad():
        inputs = processor(images=img, return_tensors="pt").to(device)
        test_features = model.get_image_features(**inputs)
        test_features = test_features / test_features.norm(p=2, dim=-1, keepdim=True)
        
        similarity = (test_features @ normal_features.T).item()
        
    if similarity < 0.8:  # threshold example
        result = "Anomaly Detected"
    else:
        result = "Normal"
    return f"Similarity: {similarity:.2f} | {result}"

gr.Interface(fn=detect_anomaly,
             inputs=gr.Image(),
             outputs="text").launch()