import gradio as gr
from PIL import Image
import torch
import numpy as np
from torchvision import transforms
from pytorch_grad_cam import GradCAM
from pytorch_grad_cam.utils.image import show_cam_on_image
import timm

device = torch.device("cpu")

model = timm.create_model("efficientnet_b4", pretrained=False, num_classes=2)
model.load_state_dict(torch.load("best_model.pth", map_location=device))
model.eval()

transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406],
                         [0.229, 0.224, 0.225])
])

mean = np.array([0.485, 0.456, 0.406])
std  = np.array([0.229, 0.224, 0.225])
target_layers = [model.conv_head]
cam = GradCAM(model=model, target_layers=target_layers)

def predict(image):
    img_tensor = transform(image).unsqueeze(0).to(device)
    with torch.no_grad():
        output = model(img_tensor)
        probs = torch.softmax(output, dim=1)[0]
        pred = output.argmax(1).item()
    fake_conf = probs[0].item()
    real_conf = probs[1].item()
    grayscale_cam = cam(input_tensor=img_tensor)
    img_np = np.array(image.resize((224, 224))).astype(np.float32) / 255.0
    cam_image = show_cam_on_image(img_np, grayscale_cam[0], use_rgb=True)
    label = "🔴 FAKE" if pred == 0 else "🟢 REAL"
    confidences = {"FAKE": round(fake_conf, 4), "REAL": round(real_conf, 4)}
    return label, confidences, Image.fromarray(cam_image)

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload a face image"),
    outputs=[
        gr.Text(label="Prediction"),
        gr.Label(label="Confidence scores"),
        gr.Image(label="Grad-CAM — What the model looks at")
    ],
    title="Deepfake Face Detector",
    description="Upload a face image to detect if its AI-generated. Model: EfficientNet-B4 trained on 140K images — 99% accuracy.",
    theme=gr.themes.Soft()
)

demo.launch()