# app.py

import torch
import torchvision.models as models
from fastai.vision.all import *
import gradio as gr

# Define the path to your model
model_path = "model.pth"

# Load the model architecture
model = models.resnet34(pretrained=False)  # Match the architecture you used
model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))  # Load weights
model.eval()  # Set to evaluation mode

# Define a simple transform function, adjust if necessary
def transform(img):
    img = PILImage.create(img).resize((224, 224))  # Resize if needed for your model
    return tensor(img).unsqueeze(0)  # Convert to tensor and add batch dimension

# Define the prediction function
def predict_image(img):
    img_tensor = transform(img)
    outputs = model(img_tensor)
    _, pred_idx = outputs.max(1)
    idx_to_class = {0: "Bike", 1: "Car", 2: "Cat", 3: "Dog"}  # Map indices to your labels
    return idx_to_class[pred_idx.item()]

# Create Gradio Interface
iface = gr.Interface(fn=predict_image, inputs="image", outputs="label", description="Upload an image to classify!")
iface.launch()