File size: 2,723 Bytes
92bcc69 3407f20 92bcc69 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 |
import gradio as gr
from PIL import Image
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import torchvision.transforms as T
class_names = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']
device = torch.device("cpu")
inference_transform = T.Compose([
T.ToTensor(),
T.Normalize(mean=(0.4914, 0.4822, 0.4465),
std=(0.2023, 0.1994, 0.2010)),
])
class SmallCifarCNN(nn.Module):
def __init__(self, num_classes: int = 10):
super().__init__()
self.features = nn.Sequential(
# Block 1
nn.Conv2d(3, 32, kernel_size=3, padding=1),
nn.BatchNorm2d(32),
nn.ReLU(inplace=True),
nn.MaxPool2d(2), # 32 -> 16
# Block 2
nn.Conv2d(32, 64, kernel_size=3, padding=1),
nn.BatchNorm2d(64),
nn.ReLU(inplace=True),
nn.MaxPool2d(2), # 16 -> 8
# Block 3
nn.Conv2d(64, 128, kernel_size=3, padding=1),
nn.BatchNorm2d(128),
nn.ReLU(inplace=True),
nn.MaxPool2d(2), # 8 -> 4
)
self.classifier = nn.Sequential(
nn.Flatten(),
nn.Linear(128 * 4 * 4, 256),
nn.ReLU(inplace=True),
nn.Dropout(p=0.5),
nn.Linear(256, num_classes),
)
def forward(self, x):
x = self.features(x)
x = self.classifier(x)
return x
deployed_model = SmallCifarCNN(num_classes=len(class_names)).to(device)
model_path = 'cifar_cnn_best.pt'
deployed_model.load_state_dict(
torch.load(model_path, map_location=device)
)
deployed_model.to(device)
deployed_model.eval()
def predict_cifar_image(img: Image.Image):
"""
Gradio callback:
- Takes a PIL Image
- Resizes to 32x32 (CIFAR size)
- Normalizes and runs through the CNN
- Returns top-3 class probabilities
"""
img = img.convert("RGB")
img = img.resize((32, 32), Image.BILINEAR)
x = inference_transform(img).unsqueeze(0).to(device)
with torch.no_grad():
logits = deployed_model(x)
probs = F.softmax(logits, dim=1).cpu().numpy().ravel()
topk = 3
idxs = np.argsort(-probs)[:topk]
return {class_names[i]: float(probs[i]) for i in idxs}
demo = gr.Interface(
fn=predict_cifar_image,
inputs=gr.Image(type="pil", label="Upload an RGB image (will be resized to 32×32)"),
outputs=gr.Label(num_top_classes=3, label="Top-3 CIFAR-10 predictions"),
title="CIFAR-10 CNN Classifier",
description="Small CNN trained on CIFAR-10. Upload an image and see top-3 class probabilities.",
)
if __name__ == "__main__":
demo.launch() |