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

print("Loading Classification Model...")
model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
model.eval()
print("Model Loaded successfully!")


LABELS = [
    "an official academic certificate or diploma",
    "a government ID card, passport, or driving license",
    "a random photograph of a person, nature, or objects",
    "a screenshot of a website or meme"
]

def verify_document_class(image: Image.Image):
    """
    Returns:
      - is_valid (bool)
      - message (str)
      - best_label (str)
      - confidence (%) (float)
      - probabilities table (list)
    """
    if image is None:
        return False, "No image provided", "", 0.0, []

    image = image.convert("RGB")

    inputs = processor(text=LABELS, images=image, return_tensors="pt", padding=True)

    with torch.no_grad():
        outputs = model(**inputs)
        probs = outputs.logits_per_image.softmax(dim=1)[0].cpu().numpy()

    best_idx = int(probs.argmax())
    best_label = LABELS[best_idx]
    confidence = float(probs[best_idx] * 100.0)

    # Gatekeeper rule
    is_valid = best_idx in (0, 1)
    message = "Valid Document Class" if is_valid else "Invalid image type. Please upload a certificate or ID."

    # Build a simple table for display
    table = [{"label": LABELS[i], "probability_%": float(probs[i] * 100.0)} for i in range(len(LABELS))]

    return is_valid, message, best_label, confidence, table


demo = gr.Interface(
    fn=verify_document_class,
    inputs=gr.Image(type="pil", label="Upload Image"),
    outputs=[
        gr.Checkbox(label="Is Valid Document?"),
        gr.Textbox(label="Message"),
        gr.Textbox(label="Best Match Label"),
        gr.Number(label="Confidence (%)"),
        gr.JSON(label="All Probabilities"),
    ],
    title="Document Class Verifier (CLIP Gatekeeper)",
    description="Uploads an image and checks whether it looks like a certificate/diploma or ID vs random photo/screenshot."
)

demo.launch(server_name="0.0.0.0", server_port=7860)