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

# 1. Load your model (Ensure this matches your training architecture)
# Change 'models.resnet18' if you used a different one
# --- UPDATED MODEL ARCHITECTURE ---
import torch.nn as nn
from torchvision import models

# 1. Initialize ResNet-50 (matches the 2048 feature size in your error)
model = models.resnet50()

# 2. Recreate the EXACT Sequential head used during your training
# This fixes the "Missing key: fc.0.weight" and "fc.3.weight" errors
model.fc = nn.Sequential(
    nn.Linear(2048, 256),   # fc.0
    nn.ReLU(),              # fc.1
    nn.Dropout(0.4),        # fc.2
    nn.Linear(256, 2)       # fc.3
)

# 3. Load your weights
model.load_state_dict(torch.load("fine_tuned_model.pt", map_location="cpu"))
model.eval()

# 2. Define labels based on your dataset folders
labels = ["Defect", "Normal"]

def predict(img):
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
    ])
    img = transform(img).unsqueeze(0)
    with torch.no_grad():
        prediction = torch.nn.functional.softmax(model(img)[0], dim=0)
        confidences = {labels[i]: float(prediction[i]) for i in range(2)}
    return confidences

# 3. Create the Interface
interface = gr.Interface(
    fn=predict, 
    inputs=gr.Image(type="pil"), 
    outputs=gr.Label(num_top_classes=2),
    title="Wall/Floor Tile Defect Inspector"
)

interface.launch()