import torch
import torch.nn.functional as F
import gradio as gr
import numpy as np
from PIL import Image
from model import CNN

# Load model
model = CNN()
model.load_state_dict(torch.load("pytorch_model.bin", map_location="cpu"))
model.eval()

# Prediction function
def predict_digit(image):
    if image is None:
        return "No image"
    image = Image.fromarray(image).convert("L").resize((28, 28))
    image = np.array(image) / 255.0
    image = torch.tensor(image).unsqueeze(0).unsqueeze(0).float()
    with torch.no_grad():
        output = model(image)
        probabilities = F.softmax(output, dim=1).numpy().flatten()
        return {str(i): float(probabilities[i]) for i in range(10)}

# Interface (no 'tool', 'type', or other unsupported args)
gr.Interface(
    fn=predict_digit,
    inputs=gr.Image(label="Upload a digit image"),
    outputs=gr.Label(num_top_classes=3),
    title="Digit Classifier",
    description="Upload a 28x28 grayscale image of a handwritten digit (0–9)."
).launch()