import torch
import gradio as gr
from torchvision import transforms
from model import HybridDeepfakeDetector

model = HybridDeepfakeDetector()
model.load_state_dict(
    torch.load("deepfake_detector_phase2.pth", 
               map_location="cpu", 
               weights_only=True)
)
model.eval()

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

def predict(image):
    tensor = transform(image).unsqueeze(0)
    with torch.no_grad():
        prob = model(tensor).item()
    print(f"Raw prob: {prob:.4f}")
    label = "REAL" if prob > 0.5 else "FAKE"
    confidence = prob if label == "REAL" else 1 - prob
    return f"{label} ({confidence*100:.1f}% confident)"

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=gr.Text(label="Prediction"),
    title="Deepfake Detector",
    description="Upload a face image to detect if it is real or AI-generated."
)

demo.launch()