A folder of your own photos on one end, a public web demo anyone can click on the other. This image classification project covers that whole distance: fine-tune a pretrained network in PyTorch, push past 90% validation accuracy, read the confusion matrix honestly, then deploy free to Hugging Face Spaces with a proper model card. Everything runs on a plain laptop CPU, no GPU and no paid account.

“A model that only lives in your notebook does not exist to anyone but you. Shipping is a feature.”

Machine learning team wisdom

Last Updated: July 2026 | Tested on: Python 3.14.6, PyTorch 2.12.1, torchvision 0.27.1, Gradio 6.20.0 | Difficulty: Advanced | Reading Time: 21 minutes

Building an image classifier is a lot like teaching a new kitchen helper to sort a delivery of vegetables. You do not teach them what an edge or a shadow is, they already see the world. You just show them a few crates labelled tomato, cucumber, and brinjal, let them make mistakes, correct the mistakes, and soon they sort the crate without you. A model pretrained on millions of photos already sees the world. Your job in this capstone is only the last mile: the labels, the honest testing, and getting the thing in front of real users.

We classify three vegetables so every block runs in seconds on any machine. The exact three do not matter. Swap in your own folders of cats and dogs, healthy and diseased leaves, or defective and clean parts, and every line below works unchanged. The dataset is the excuse. The pipeline is the lesson.

What We Are Building

By the end you have a folder that runs top to bottom and produces three things a hiring manager actually clicks: a trained model file, an evaluation you can defend, and a public URL where anyone can drop in a photo and get a prediction. The whole thing follows one repeatable shape you can reuse for any image task at work or in a portfolio. The diagram below is that shape in one picture.

One honest note up front. In a real project you collect photos with a camera or scrape them with permission. So that every block here always runs and can never break because a download link died, we draw a small stand-in dataset in code with three clearly different vegetable shapes. The moment you point the same code at real image folders, nothing changes except the pictures.

Prerequisites

The Capstone Path

Read the order carefully, because in an image classification project the order is the whole point. You gather labelled images, borrow a network that already knows how to see, train only a small new head, and then stop to check the accuracy gate. Above 90% you move on to the honest audit and shipping. Below it you loop back and keep tuning. Notice that evaluation, the model card, and deployment come after training, not instead of it. A model nobody can use and nobody can trust is not finished, no matter how good the accuracy looks.

Step 1: Assemble a Small Custom Dataset

PyTorch has a beautifully simple convention for image data: one folder per class, with all of that class’s images inside it, split into a train and a val directory. Get your files into that shape and datasets.ImageFolder reads the labels straight from the folder names. Here we draw the images so the post is self-contained, adding rotation and noise so the classifier has to work a little rather than memorize perfect clip art.

📄 make_dataset.py: draw a small 3-class dataset in the folders-per-label layout

import numpy as np
from PIL import Image, ImageDraw
from pathlib import Path

rng = np.random.default_rng(42)
CLASSES = ["tomato", "cucumber", "brinjal"]
ROOT, SIZE = Path("veggies"), 224

def draw_one(kind):
    bg = rng.integers(215, 245)
    img = Image.new("RGB", (SIZE, SIZE), (bg, bg, bg))
    d = ImageDraw.Draw(img)
    cx, cy = rng.integers(80, 144, 2)
    if kind == "tomato":                       # round and red
        r = rng.integers(30, 44)
        d.ellipse([cx-r, cy-r, cx+r, cy+r],
                  fill=(rng.integers(150, 220), rng.integers(30, 90), 40))
    elif kind == "cucumber":                   # slim and green
        w, h = rng.integers(14, 22), rng.integers(44, 62)
        d.ellipse([cx-w, cy-h, cx+w, cy+h], fill=(40, rng.integers(90, 150), 40))
    else:                                       # brinjal: slim and purple
        w, h = rng.integers(18, 27), rng.integers(36, 52)
        d.ellipse([cx-w, cy-h, cx+w, cy+h],
                  fill=(rng.integers(120, 195), 28, rng.integers(110, 165)))
    img = img.rotate(rng.integers(-35, 35), fillcolor=(bg, bg, bg))
    noisy = np.clip(np.asarray(img, float) + rng.normal(0, 26, (SIZE, SIZE, 3)), 0, 255)
    return Image.fromarray(noisy.astype("uint8"))

for split, n in [("train", 80), ("val", 30)]:
    for c in CLASSES:
        (ROOT / split / c).mkdir(parents=True, exist_ok=True)
        for i in range(n):
            draw_one(c).save(ROOT / split / c / f"{c}_{i:03d}.png")
    print(f"{split}: {n * len(CLASSES)} images ({', '.join(c + '=' + str(n) for c in CLASSES)})")

▶ Output

train: 240 images (tomato=80, cucumber=80, brinjal=80)
val: 90 images (tomato=30, cucumber=30, brinjal=30)

What happened here: We now have 240 training images and 90 validation images (30 per class, the last line prints the per-class count for the split) sitting in veggies/train/<class>/ and veggies/val/<class>/. Two of the three classes, cucumber and brinjal, are both slim vertical shapes, and after a hard rotation and a dose of noise they start to look alike. That is on purpose. A dataset where every class is trivially separable teaches you nothing about reading a confusion matrix, and real photos are never that clean.

Step 2: Transfer-Learn a torchvision Model

Now the borrowed brain. We load a ResNet-18 that already learned to see on ImageNet, freeze its whole body so those hard-won features stay put, and replace only the final layer with a fresh one sized to our three classes. Then we run the standard training loop: pull batches from the DataLoader, compute the loss, backpropagate, step the optimizer. Because only the tiny new head learns, this converges in a handful of epochs even on a CPU.

One detail that quietly decides whether this works: the input preprocessing must match what the pretrained weights expect. That means resizing to 224 pixels and normalizing with the exact ImageNet mean and standard deviation. We pin the weights explicitly with ResNet18_Weights.IMAGENET1K_V1 rather than DEFAULT, so an upgrade of torchvision cannot silently swap the weights under you and change your numbers.

📄 train.py: freeze the ResNet-18 backbone, train a new head, print the curve

import torch, torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms, models

torch.manual_seed(42)
device = "cuda" if torch.cuda.is_available() else "cpu"

weights = models.ResNet18_Weights.IMAGENET1K_V1   # pinned, not DEFAULT
mean, std = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
train_tf = transforms.Compose([
    transforms.RandomResizedCrop(224, scale=(0.8, 1.0)),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(), transforms.Normalize(mean, std)])
eval_tf = transforms.Compose([
    transforms.Resize(224), transforms.CenterCrop(224),
    transforms.ToTensor(), transforms.Normalize(mean, std)])

train_ds = datasets.ImageFolder("veggies/train", transform=train_tf)
val_ds = datasets.ImageFolder("veggies/val", transform=eval_tf)
classes = train_ds.classes
train_dl = DataLoader(train_ds, batch_size=32, shuffle=True)
val_dl = DataLoader(val_ds, batch_size=32)

model = models.resnet18(weights=weights)
for p in model.parameters():
    p.requires_grad = False                       # freeze the ImageNet body
model.fc = nn.Linear(model.fc.in_features, len(classes))   # only this head trains
model = model.to(device)

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.fc.parameters(), lr=1e-3)

def run(dl, train):
    model.train() if train else model.eval()
    loss_sum = correct = total = 0
    for x, y in dl:
        x, y = x.to(device), y.to(device)
        with torch.set_grad_enabled(train):
            out = model(x)
            loss = criterion(out, y)
            if train:
                optimizer.zero_grad(); loss.backward(); optimizer.step()
        loss_sum += loss.item() * x.size(0)
        correct += (out.argmax(1) == y).sum().item()
        total += x.size(0)
    return loss_sum / total, correct / total

print(f"Classes: {classes}  |  device: {device}")
print(f"{'epoch':>5} {'train_loss':>11} {'val_loss':>9} {'val_acc':>8}")
for epoch in range(1, 7):
    tl, _ = run(train_dl, True)
    vl, va = run(val_dl, False)
    print(f"{epoch:>5} {tl:>11.3f} {vl:>9.3f} {va:>8.1%}")

torch.save({"state_dict": model.state_dict(), "classes": classes}, "veggie_resnet18.pt")
print("Saved veggie_resnet18.pt")

▶ Output

Classes: ['brinjal', 'cucumber', 'tomato']  |  device: cpu
epoch  train_loss  val_loss  val_acc
    1       0.996     0.893    74.4%
    2       0.670     0.634    87.8%
    3       0.470     0.457    92.2%
    4       0.341     0.374    92.2%
    5       0.283     0.312    95.6%
    6       0.235     0.289    95.6%
Saved veggie_resnet18.pt

What happened here: Both losses fall every epoch and validation accuracy climbs from 74.4% to 95.6%, comfortably clearing the 90% gate the diagram asked for. Because training loss and validation loss drop together, the model is genuinely learning rather than memorizing the training set, which would show up as validation loss creeping back up. On a CPU this whole run takes a couple of minutes; on a GPU it is seconds. We saved the trained weights and the class list into one .pt file so the next scripts, and the live demo, can load exactly what we trained without guessing the label order.

Step 3: Evaluate Honestly

A single accuracy number is a headline, not a story. 95.6% could hide the fact that one class is a disaster while the other two are perfect, and that average would never tell you. The confusion matrix breaks the score open: each row is the true class, each column is what the model guessed, so the diagonal is correct and everything off the diagonal is a specific mistake you can go look at. Say a reviewer named Aditi asks “which vegetable does it get wrong, and why.” This is how you answer her.

📄 evaluate.py: confusion matrix, per-class report, and the list of misses

import torch, torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms, models
from sklearn.metrics import confusion_matrix, classification_report

ckpt = torch.load("veggie_resnet18.pt", weights_only=False)
classes = ckpt["classes"]
model = models.resnet18(weights=None)
model.fc = nn.Linear(model.fc.in_features, len(classes))
model.load_state_dict(ckpt["state_dict"]); model.eval()

eval_tf = transforms.Compose([
    transforms.Resize(224), transforms.CenterCrop(224), transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
val_ds = datasets.ImageFolder("veggies/val", transform=eval_tf)
val_dl = DataLoader(val_ds, batch_size=32)

y_true, y_pred, misses, idx = [], [], [], 0
with torch.no_grad():
    for x, y in val_dl:
        p = model(x).argmax(1)
        y_true += y.tolist(); y_pred += p.tolist()
        for yi, pi in zip(y.tolist(), p.tolist()):
            name = val_ds.samples[idx][0].replace("\\", "/").split("/")[-1]
            if yi != pi:
                misses.append((name, classes[yi], classes[pi]))
            idx += 1

cm = confusion_matrix(y_true, y_pred)
print("Confusion matrix (rows = true, cols = predicted)")
print("           " + "".join(f"{c[:8]:>10}" for c in classes))
for i, c in enumerate(classes):
    print(f"{c[:10]:>10} " + "".join(f"{n:>10}" for n in cm[i]))
print("\n" + classification_report(y_true, y_pred, target_names=classes, digits=3))
print(f"Misclassified {len(misses)} of {len(y_true)} validation images:")
for name, true, pred in misses:
    print(f"  {name:<16} true={true:<9} predicted={pred}")

▶ Output

Confusion matrix (rows = true, cols = predicted)
              brinjal  cucumber    tomato
   brinjal         30         0         0
  cucumber          4        26         0
    tomato          0         0        30

              precision    recall  f1-score   support

     brinjal      0.882     1.000     0.938        30
    cucumber      1.000     0.867     0.929        30
      tomato      1.000     1.000     1.000        30

    accuracy                          0.956        90
   macro avg      0.961     0.956     0.955        90
weighted avg      0.961     0.956     0.955        90

Misclassified 4 of 90 validation images:
  cucumber_000.png true=cucumber  predicted=brinjal
  cucumber_002.png true=cucumber  predicted=brinjal
  cucumber_014.png true=cucumber  predicted=brinjal
  cucumber_017.png true=cucumber  predicted=brinjal

What happened here: Now the 95.6% has a shape. Tomato is perfect, every one of its 30 images was correct, because a red round blob looks nothing like the two slim green and purple shapes. Every single mistake sits in one cell: four cucumbers predicted as brinjal, and never the reverse. That is the misclassification gallery, and it is worth looking at rather than hiding. Those four are slim vertical shapes that a hard rotation and heavy noise pushed toward brinjal’s proportions.

The model even tells you it was unsure, not confidently wrong, which we can prove in the next step. Reading precision and recall confirms it: cucumber recall is 0.867 because four of its thirty were missed, while brinjal precision is 0.882 because it absorbed those four false positives.

Step 4: Wrap It in a Gradio Interface

An image classification project that ends at a trained .pt file is useless to anyone who cannot write Python. Gradio fixes that in about ten lines: you give it a function that takes an image and returns predictions, and it builds a web page with a drag-and-drop upload box and a nice bar chart of the results. The heart of it is the classify function, which does exactly what evaluation did for one image and returns a label-to-probability dictionary that Gradio renders as a chart.

📄 app.py: the Gradio demo (also the file Hugging Face Spaces runs)

import torch, torch.nn as nn, gradio as gr
from torchvision import transforms, models
from PIL import Image

ckpt = torch.load("veggie_resnet18.pt", weights_only=False)
classes = ckpt["classes"]
model = models.resnet18(weights=None)
model.fc = nn.Linear(model.fc.in_features, len(classes))
model.load_state_dict(ckpt["state_dict"]); model.eval()

tf = transforms.Compose([
    transforms.Resize(224), transforms.CenterCrop(224), transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])

def classify(img: Image.Image) -> dict:
    x = tf(img.convert("RGB")).unsqueeze(0)
    with torch.no_grad():
        probs = torch.softmax(model(x), dim=1)[0]
    return {classes[i]: float(probs[i]) for i in range(len(classes))}

demo = gr.Interface(
    fn=classify,
    inputs=gr.Image(type="pil", label="Upload a vegetable photo"),
    outputs=gr.Label(num_top_classes=3, label="Prediction"),
    title="Veggie Classifier",
    description="Fine-tuned ResNet-18. Predicts tomato, cucumber, or brinjal.")

if __name__ == "__main__":
    demo.launch()

Before launching a server, it is worth testing the classify function on its own, the way a good test would. We run it on one clean tomato and on one of the four cucumbers the confusion matrix flagged, and print the probabilities.

▶ Output (classify() called directly on two validation images)

tomato_005.png -> predicted tomato
    tomato     82.4%
    brinjal    13.1%
    cucumber    4.6%
cucumber_000.png -> predicted brinjal
    brinjal    55.0%
    cucumber   33.5%
    tomato     11.5%

What happened here: The tomato is an easy, confident call at 82.4%. The interesting one is cucumber_000.png, a miss from the confusion matrix: the model picked brinjal at 55.0% but put cucumber right behind at 33.5%. That is a close, honest kind of wrong, not a wild guess, and it matches the story the confusion matrix told. When you call demo.launch() locally, Gradio starts a small web server and prints a URL you open in a browser to try the drag-and-drop app yourself.

▶ Example output (from demo.launch(), running locally)

* Running on local URL:  http://127.0.0.1:7860
To create a public link, set `share=True` in `launch()`.

Step 5: Deploy to Hugging Face Spaces

A local URL only works while your laptop is on and only for you. To hand out a link that lives on the public internet, we push the app to a host. At the time of writing, Hugging Face Spaces is the standard free home for a Gradio demo like this: it gives you a Git repository, installs your dependencies, runs app.py, and serves it at a public URL at no cost. If Spaces ever changes its terms or you outgrow the free tier, Streamlit Community Cloud and Modal are the two common fallbacks, and your app.py plus model file move across almost unchanged.

A Space needs exactly three files next to your model: the app.py from above, a requirements.txt so the host installs the right libraries, and a README.md whose header tells Spaces this is a Gradio app. Create a new Space (choose the Gradio SDK), then either drag these files into the web uploader or push them with Git.

📄 requirements.txt: pin the libraries so the Space builds the same every time

torch==2.12.1
torchvision==0.27.1
gradio==6.20.0
pillow==12.2.0

📄 README.md: the header block Hugging Face Spaces reads to configure the app

---
title: Veggie Classifier
emoji: 🍅
sdk: gradio
sdk_version: 6.20.0
app_file: app.py
---

# Veggie Classifier
A fine-tuned ResNet-18 that sorts photos into tomato, cucumber, or brinjal.

Push those four files (the two above, plus app.py and veggie_resnet18.pt) and the Space builds itself. After a minute or two of installing, the status turns to Running and you get a shareable address of the form https://huggingface.co/spaces/your-name/veggie-classifier. That link is the deliverable: send it to anyone and they can upload a photo and see the model work, no install and no code.

Step 6: Write the Model Card

A model card is the label on the tin. It is a short honest document that says what the model was trained on, what it is meant to do, and where it falls down, so nobody uses it for something it was never built for. Skipping it is how a toy trained on three cartoon vegetables ends up quietly deployed in a grocery app and failing on real produce. On Hugging Face the card is just the rest of that same README.md, in plain Markdown.

📄 The model card appended to README.md: data, intended use, and limits

## Model details
- Base model: torchvision ResNet-18, ImageNet1K_V1 weights, backbone frozen.
- Head: one linear layer over 3 classes, trained for 6 epochs on CPU.
- Reported result: 95.6% accuracy on a held-out validation set of 90 images.

## Intended use
- A teaching demo for an end-to-end image classification project.
- Sorting images into exactly three shapes: tomato, cucumber, brinjal.

## Limitations and biases
- Trained on synthetic drawn shapes, NOT real photographs. It will not
  generalize to real produce, other vegetables, or multiple items per image.
- Weakest on cucumber, which is sometimes confused with brinjal (both slim);
  see the confusion matrix. Do not use for any decision that matters.

## How to use
Upload a single image. The app returns a probability for each of the 3 classes.

What happened here: Notice the card leads with limitations, not with the accuracy. That is the tell of a card written by someone who has watched a model get misused. The 95.6% is real, but it was measured on drawn shapes, so the card says loudly that this will not work on real produce. Written this way, the card protects both the next engineer and you, because it records exactly what you claimed and did not claim.

Common Mistakes

Mistake 1: Different preprocessing at train and predict time

If you train with ImageNet normalization but forget it in the Gradio app, or resize to a different size, the model sees inputs it was never trained on and accuracy quietly collapses on the live demo while looking fine in your notebook. Define the eval transform once and use the exact same one in evaluate.py and app.py. A silent preprocessing mismatch is the single most common reason a deployed image model underperforms its reported score.

Mistake 2: Trusting one accuracy number

An overall 95.6% felt great until the confusion matrix showed every mistake landed on cucumber. On an imbalanced or safety-sensitive task that blind spot matters enormously. Always print the confusion matrix and the per-class report, and actually open the misclassified images. The average hides exactly the failure you most need to see.

Mistake 3: Shipping without a model card

A demo with no stated limits invites misuse. Someone finds your Space, assumes it works on real photos, and wires it into something that matters. Two paragraphs of honest scope and limitations cost you five minutes and save the next person hours of confusion. A public model without a card is an unlabelled bottle in the medicine cabinet.

Best Practices

• DO start with feature extraction (freeze the backbone, train only the head) and reach for full fine-tuning only if accuracy stalls below your gate
• DO pin the pretrained weights by name (ResNet18_Weights.IMAGENET1K_V1) so a library upgrade cannot change your results underneath you
• DO save the class list alongside the weights, so the demo never guesses the label order
• DO judge the model by its confusion matrix and misclassified images, not by a single accuracy figure
• DON’T let train-time and predict-time preprocessing drift apart; share one transform
• DON’T publish a public demo without a model card that states the data, the intended use, and the limits

The core here (torchvision models, the PyTorch training loop, a Gradio interface) has been stable for years and will outlive any single hosting service. At the time of writing Hugging Face Spaces is the easy free host, but if that changes, the same app.py and model file deploy to Streamlit Community Cloud or Modal with only the wrapper swapped.

Conclusion

You ran a full deep learning image classification project end to end. You assembled a small custom dataset in the folders-per-label layout, transfer-learned a frozen ResNet-18 with the standard training loop and cleared 90% validation accuracy, then audited the result honestly with a confusion matrix that pinned every mistake on one class. You wrapped the model in a Gradio interface, deployed it free to Hugging Face Spaces as a public link, and wrote a model card that leads with its limits. That last mile, from a good number to a trustworthy shared demo, is exactly what separates a course exercise from a portfolio piece.

Point this same pipeline at your own image folders this week and deploy the result. For the full path from Python basics through deep learning and deployment, browse the Python + AI/ML tutorial series home.

Frequently Asked Questions

Interview Questions on Deep Learning Projects

Interviewers rarely ask for definitions. They ask what happens in situations like these.

Q: Walk me through an image classification project end to end.

Gather labelled images in a folders-per-class layout and split into train and validation. Load a pretrained network like ResNet-18, freeze the backbone, and replace the final layer with one sized to your classes. Train that head with a standard loop until validation accuracy clears your target, then audit with a confusion matrix and by opening the misclassified images rather than trusting one number. Finally wrap the model in an interface, deploy it to a public host, and write a model card stating the data, intended use, and limits.

Q: Your model shows 96% accuracy but the confusion matrix worries you. Why might that be?

A high average can hide a class that fails badly. In this project every mistake fell on cucumber, which was confused with brinjal, while the other two classes were near perfect. If cucumber were the safety-critical class, a 96% average would be dangerously misleading. That is why you read per-class recall and precision and inspect the actual misclassified images, so the average never lets a specific failure hide.

Q: A deployed image classifier performs worse in production than in your notebook. What do you check first?

Preprocessing first. The most common cause is that the live app resizes or normalizes inputs differently from training, so the model sees a distribution it never learned. I would confirm the same input size and the same ImageNet mean and standard deviation are applied in both places, and that the class-index order matches. After that I would look at genuine distribution shift, meaning real inputs that differ from the training images, which the model card should already have warned about.

Q: Why freeze the backbone instead of training the whole network?

The pretrained backbone already converts pixels into rich general features from millions of ImageNet images. Freezing it means only a small head learns, which trains in minutes on a CPU, needs far less data, and resists overfitting a small dataset. You unfreeze and fully fine-tune only when your images differ a lot from ImageNet or feature extraction stalls below your accuracy target, and then you use a much smaller learning rate so you do not wipe out the borrowed features.

Q: What belongs in a model card and why does it matter for deployment?

The training data and its nature, the intended use, the measured performance, and the limitations and known biases. It matters because a public model with no stated scope gets used for things it was never built for. A card that leads with limits, for example stating this was trained on synthetic shapes and will not generalize to real produce, protects the next user and records exactly what you did and did not claim.

Q: How would you take this demo from a toy to production-ready?

Replace the synthetic data with a large, well-labelled set of real photos and re-audit the confusion matrix on genuinely unseen images. Add input validation and confidence thresholds so low-confidence predictions are flagged rather than guessed. Move off the free tier to a container you control, add monitoring for drift so you catch the model degrading over time, and keep the model card and versioned weights updated as the data changes.

Related Posts

Previous: DL: RNNs and LSTMs for Sequence Processing

Next: The Attention Mechanism From Scratch in PyTorch

Series Home: Python + AI/ML Tutorial Series