AaronJames95/MNIST_Digit_Classifier

1import gradio as gr2from fastai.vision.all import *3from PIL import Image, ImageOps4import numpy as np5 6# Load your trained model7learn = load_learner('mnist_resnet18.pkl')8 9def predict_from_sketch(img):10    img = img["composite"]11 12    # Convert to grayscale, invert, resize, RGB13    pil_img = Image.fromarray(img).convert("RGBA")14 15    # Composite over a white background16    background = Image.new("RGBA", pil_img.size, (255, 255, 255, 255))17    pil_img = Image.alpha_composite(background, pil_img)18    19    # Now convert to grayscale properly20    pil_img = pil_img.convert("L")         21 22    # Invert and resize23    pil_img = ImageOps.invert(pil_img)24    pil_img = pil_img.resize((224, 224)).convert("RGB")25    #pil_img.show(title="Final Preprocessed Image")26 27    dl = learn.dls.test_dl([pil_img])28    xb = dl.one_batch()[0]29    30 31    with torch.no_grad():32        preds = learn.model.eval()(xb)33        pred_idx = preds.argmax(dim=1).item()34        probs = preds.softmax(dim=1).squeeze()35    return {str(learn.dls.vocab[i]): float(probs[i]) for i in range(len(probs))}36 37demo = gr.Interface(38    predict_from_sketch,39    inputs="sketchpad",40    outputs=gr.Label(num_top_classes=3),41    title="MNIST Digit Classifier",42    description="Draw a digit (0–9) and get predictions in real time!"43)44 45if __name__ == "__main__":46    demo.launch()47