johnphs-9/6m-data-3.8-Computer-Vision

1# app.py — Fashion-MNIST clothing recogniser (Hugging Face Space)2# ----------------------------------------------------------------3# This is the file Hugging Face runs to build your web app.4# It loads the model YOU trained (tiny_cnn.pt) and puts an5# "upload a photo -> get a prediction" page in front of it.6 7import torch8import torch.nn as nn9import torchvision.transforms as T10import gradio as gr11 12# ----------------------------------------------------------------13# 1) The model definition.14#    This MUST be the exact same TinyCNN from your training notebook,15#    otherwise the saved weights won't fit. (Copied from 03_first_cnn.)16# ----------------------------------------------------------------17class TinyCNN(nn.Module):18    def __init__(self, n_classes=10):19        super().__init__()20        self.features = nn.Sequential(21            nn.Conv2d(1, 16, kernel_size=3, padding=1),22            nn.ReLU(),23            nn.MaxPool2d(2),24            nn.Conv2d(16, 32, kernel_size=3, padding=1),25            nn.ReLU(),26            nn.MaxPool2d(2),27        )28        self.classifier = nn.Sequential(29            nn.Flatten(),30            nn.Linear(32 * 7 * 7, 64),31            nn.ReLU(),32            nn.Linear(64, n_classes),33        )34 35    def forward(self, x):36        return self.classifier(self.features(x))37 38 39# The 10 classes Fashion-MNIST knows about — IN THIS ORDER (0..9).40CLASS_NAMES = [41    "T-shirt/top", "Trouser", "Pullover", "Dress", "Coat",42    "Sandal", "Shirt", "Sneaker", "Bag", "Ankle boot",43]44 45# ----------------------------------------------------------------46# 2) Load the trained weights once, when the app starts up.47# ----------------------------------------------------------------48model = TinyCNN()49model.load_state_dict(torch.load("tiny_cnn.pt", map_location="cpu"))50model.eval()  # eval mode = "we're predicting, not training"51 52# ----------------------------------------------------------------53# 3) Turn ANY uploaded photo into what the model expects:54#    a 28x28 GRAYSCALE image.55#    Fashion-MNIST images are light clothing on a BLACK background,56#    so we invert (most phone photos are dark item on light background).57# ----------------------------------------------------------------58preprocess = T.Compose([59    T.Grayscale(num_output_channels=1),  # colour -> grey60    T.Resize((28, 28)),                  # shrink to 28x2861    T.functional.invert,                 # flip light/dark to match training data62    T.ToTensor(),                        # -> tensor with values 0..163])64 65 66def predict(image):67    if image is None:68        return {}69    x = preprocess(image).unsqueeze(0)   # add a batch dimension: (1, 1, 28, 28)70    with torch.no_grad():71        logits = model(x)72        probs = torch.softmax(logits, dim=1)[0]  # turn scores into probabilities73    # Gradio's Label widget wants {class_name: probability}74    return {CLASS_NAMES[i]: float(probs[i]) for i in range(len(CLASS_NAMES))}75 76 77# ----------------------------------------------------------------78# 4) Build the web page.79# ----------------------------------------------------------------80description = (81    "Upload an image and my CNN will guess which of 10 clothing categories it is.\n\n"82    "⚠️ This model only ever saw tiny 28×28 grayscale clothing images during training "83    "(Fashion-MNIST). It will confidently guess on ANY image — even ones it should not "84    "recognise. Try different photos and notice when it succeeds and when it fails!"85)86 87demo = gr.Interface(88    fn=predict,89    inputs=gr.Image(type="pil", label="Upload an image"),90    outputs=gr.Label(num_top_classes=3, label="Top guesses"),91    title="My Fashion-MNIST Clothing Recogniser",92    description=description,93    flagging_mode="never",94)95 96if __name__ == "__main__":97    demo.launch()98