app.py

import torch
import gradio as gr
from PIL import Image
import torchvision.transforms as T

from models.encoder import ResNetEncoder
from models.metaoptnet_svm import MetaOptNetSVM
from build_support_set import build_support_set

# Device setup to use CPU
device = torch.device("cpu")

# Load model
encoder = ResNetEncoder()
encoder.load_state_dict(torch.load("models/encoder_ep8.pth", map_location=device))
encoder = encoder.to(device)
encoder.eval()

classifier = MetaOptNetSVM()

# Build support set once at startup
support_x, support_y = build_support_set("data", num_per_class=5)
support_x, support_y = support_x.to(device), support_y.to(device)

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

def predict(image: Image.Image) -> str:
    query_x = transform(image.convert("RGB")).unsqueeze(0).to(device)
    support_features = encoder(support_x)
    query_features = encoder(query_x)

    preds = classifier(support_features, support_y, query_features)
    label = preds.argmax(dim=1).item()
    return "Fake" if label == 1 else "Real"

# Launch Gradio app
gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs="text",
    title="Deepfake Image Detector",
    description="Upload an image to classify it as Real or Fake using meta-learning."
).launch()