app.py

import gradio as gr
import torch
import torch.nn as nn
from torchvision import models, transforms
from safetensors.torch import load_file
from PIL import Image

# 1. Rebuild the blank ResNet-50 architecture
model = models.resnet50(weights=None)
model.fc = nn.Linear(model.fc.in_features, 2) # 2 classes: Fake (0) and Real (1)

# 2. Load YOUR trained weights from the safetensors file
# Using map_location='cpu' ensures it works on Hugging Face's free CPU tier!
state_dict = load_file("model.safetensors", device="cpu")

# (Optional safety check) If you trained with DataParallel, keys might have "module." in front.
# This removes it so the weights load perfectly no matter what.
state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()}

# Inject the weights into the skeleton and set to test mode
model.load_state_dict(state_dict)
model.eval()

# 3. Define the strict patch transform (No randomness!)
test_transform = transforms.Compose([
    transforms.Resize(256), # Standard practice to resize slightly before center cropping
    transforms.CenterCrop(224), 
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

classes = ['FAKE', 'REAL']

# 4. Create the prediction function that Gradio will call
def predict_image(img):
    if img is None:
        return None
        
    # Ensure image is strictly RGB (drops alpha channels from PNGs)
    img = img.convert('RGB')
    
    # Apply transforms and add the batch dimension
    img_tensor = test_transform(img).unsqueeze(0) 
    
    with torch.no_grad():
        # Pass through model
        preds = model(img_tensor)
        # Convert raw numbers to percentages (0.0 to 1.0)
        probs = torch.nn.functional.softmax(preds[0], dim=0)
        
    # Gradio expects a dictionary of { "Class Name": probability }
    return {classes[i]: float(probs[i]) for i in range(2)}

# 5. Build and launch the Web App!
interface = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(type="pil", label="Upload an Image"),
    outputs=gr.Label(num_top_classes=2, label="Prediction"),
    title="PixelSleuth: AI Image Detector 🕵️‍♂️",
    description="Upload an image to see if it is a Real Photograph or AI-Generated (Fake). The model analyzes microscopic pixel artifacts to make its decision.",
    allow_flagging="never" # Turns off the confusing "Flag" button for users
)

interface.launch()