import gradio as gr
import torch
from transformers import ViTFeatureExtractor, ViTForImageClassification
from PIL import Image
import requests
import numpy as np

# Load the pretrained ViT model and feature extractor
model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224-in21k')
feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224-in21k')

# Define the prediction function
def predict_deepfake(image):
    # Preprocess the image
    inputs = feature_extractor(images=image, return_tensors="pt")

    # Make predictions
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
        predicted_class_idx = logits.argmax(-1).item()

    # Assuming you have a dictionary of labels mapping index to class names (Real vs Fake)
    class_labels = {0: "Real", 1: "Fake"}  # This is just a placeholder
    return class_labels.get(predicted_class_idx, "Unknown")

# Gradio Interface
iface = gr.Interface(
    fn=predict_deepfake,
    inputs=gr.Image(type="pil", label="Upload Image"),
    outputs=gr.Textbox(label="Prediction"),
    title="Deepfake Detector",
    description="Upload an image to classify it as 'Real' or 'Fake' using Vision Transformer (ViT).",
    allow_flagging="never"
)

# Launch the Gradio app
if __name__ == "__main__":
    iface.launch()