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
from torchvision import models
model = models.resnet18()
model.fc = torch.nn.Linear(model.fc.in_features, 2) 
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()