# app.py
import gradio as gr
from transformers import AutoImageProcessor, AutoModelForImageClassification
from PIL import Image
import torch

# Load processor and model
processor = AutoImageProcessor.from_pretrained("nguyenkhoa/dinov2_Liveness_detection_v2.2.3")
model = AutoModelForImageClassification.from_pretrained("nguyenkhoa/dinov2_Liveness_detection_v2.2.3")

# Define labels
id2label = model.config.id2label

# Inference function
def detect_liveness(image: Image.Image):
    # Preprocess image
    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)

    logits = outputs.logits
    probs = torch.nn.functional.softmax(logits, dim=-1)[0]
    
    # Get prediction
    predicted_class_idx = torch.argmax(probs).item()
    predicted_label = id2label[predicted_class_idx]
    confidence = round(probs[predicted_class_idx].item(), 4)

    return f"Liveness: {predicted_label} (Confidence: {confidence})"

# Launch Gradio app
app = gr.Interface(
    fn=detect_liveness,
    inputs=gr.Image(type="pil", label="Upload Face Image"),
    outputs=gr.Text(label="Liveness Detection Result"),
    title="Liveness Detection App",
    description="Upload a face image to check if it's live or spoofed using DinoV2 model."
)

if __name__ == "__main__":
    app.launch()