import torch
from transformers import AutoModel, AutoImageProcessor
from PIL import Image
import gradio as gr
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity

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

# Load DINOv2 model
model_name = "facebook/dinov2-base"
model = AutoModel.from_pretrained(model_name).to(device)
processor = AutoImageProcessor.from_pretrained(model_name)

# Load OK reference image (built-in normal)
#ok_image = Image.open("https://huggingface.co/datasets/Soooma/sofc/blob/main/OK1.jpg")
ok_image = Image.open("OK1.jpg")

# Precompute OK feature embedding
with torch.no_grad():
    ok_input = processor(images=ok_image, return_tensors="pt").to(device)
    ok_feat = model(**ok_input).last_hidden_state.mean(dim=1)
    ok_feat = ok_feat.cpu().numpy()

def detect_anomaly(image):
    if image is None:
        return "No image uploaded."

    with torch.no_grad():
        inputs = processor(images=image, return_tensors="pt").to(device)
        feat = model(**inputs).last_hidden_state.mean(dim=1)
        feat = feat.cpu().numpy()

    similarity = cosine_similarity(feat, ok_feat)[0][0]

    if similarity < 0.90:
        return f"Anomaly Detected | Similarity: {similarity:.3f}"
    else:
        return f"Normal | Similarity: {similarity:.3f}"

gr.Interface(
    fn=detect_anomaly,
    inputs=gr.Image(type="pil"),
    outputs="text",
    title="Anomaly Detector (DINOv2)",
    description="Upload an image of a stack. The model compares it to a known OK sample using DINOv2 features."
).launch()