Fine-grained image classification: 37 cat and dog breeds from the Oxford-IIIT Pet dataset, using an ImageNet-pretrained EfficientNet-B0 fine-tuned end to end. Every prediction ships with a Grad-CAM map showing which pixels actually drove it.
Top row: the photo. Bottom row: Grad-CAM — red is the evidence the model used. The last column is a genuine mistake, included deliberately.
docker compose up --buildThen open http://localhost:8000. If port 8000 is taken:
PORT=8010 docker compose up --build # http://localhost:8010The fine-tuned checkpoint is committed to this repository and copied into the image, so the container classifies immediately — no dataset download, no training, no network access at runtime.
Requires uv and Python 3.11+.
uv sync
uv run main.py # http://127.0.0.1:8000Drop, paste, or browse to a photo of a cat or dog — or click one of the four bundled examples, which come from the official test split. You get:
- the top-5 breeds with probabilities,
- a Grad-CAM overlay of the pixels behind the top prediction,
- a close-call warning when the top two breeds are within 25 percentage points, which on this dataset usually means one of the terrier pairs.
This is the case where training from scratch simply does not work. There are 37 classes with roughly 100 training images each — nowhere near enough to learn general visual features. ImageNet pretraining supplies those; only the head and a light fine-tune are learned here.
The schedule is deliberately two-phase:
- Head-only warm-up. The backbone is frozen while the randomly-initialised classifier learns. Skipping this lets huge early gradients from the random head flow back and damage the pretrained features.
- Full fine-tune at a much lower learning rate, with cosine decay.
Phase 1 alone reaches 85.5% validation top-1 — before a single backbone weight has moved. That is the whole argument for transfer learning in one number.
Trained for 30 epochs at 224×224 in 5.3 minutes on a single RTX 2060 SUPER, scored on the official 3669-image test split using the checkpoint with the best validation top-1.
| Metric | Value |
|---|---|
| Test top-1 | 90.9% |
| Test top-5 | 99.1% |
| Macro F1 | 0.907 |
| Best validation top-1 | 92.8% |
Top-5 at 99.1% is the more revealing number: the correct breed is essentially always in the model's shortlist, and nearly all remaining error is the final choice between two genuinely similar breeds.
Macro F1 tracks top-1 closely (0.907 vs 0.909), which says accuracy is spread fairly evenly rather than propped up by a few large classes.
Two runs, selected on validation only:
| Run | Fine-tune LR | Augmentation | Val top-1 | Test top-1 |
|---|---|---|---|---|
| 1 | 3e-4 | crops 0.6–1.0, jitter 0.25, erase 0.25 | 91.5% | 89.6% |
| 2 (shipped) | 1.5e-4 | crops 0.7–1.0, jitter 0.2, erase 0.15 | 92.8% | 90.9% |
Run 1 plateaued from epoch 5 — the signature of too high a learning rate for a small dataset, combined with augmentation aggressive enough to destroy the coat markings that distinguish several of these breeds. Halving the LR and easing the crops was worth +1.3 points of test accuracy.
| Breed | F1 |
|---|---|
| Staffordshire Bull Terrier | 0.615 |
| American Pit Bull Terrier | 0.629 |
| Ragdoll | 0.726 |
| American Bulldog | 0.770 |
| Birman | 0.784 |
These are not random errors. The three worst are the Staffordshire Bull Terrier / American Pit Bull Terrier / American Bulldog cluster — breeds so similar that shelters and breed-identification studies find humans disagree on them too. Ragdoll and Birman are the same story among cats: both are colourpoint, blue-eyed, semi-long-haired. The model is not confused about dogs versus cats; it is confused exactly where the label boundary is genuinely thin.
Rows are the true breed, columns the prediction. The teal diagonal is correct answers; red off-diagonal cells are confusions, and they cluster exactly where you would expect.
- Closed set. The model knows 37 breeds and nothing else. Show it a rabbit, a mixed-breed dog, or a sofa and it will still confidently name a breed — there is no "none of the above" class and no out-of-distribution detection.
- Single run, one seed. No cross-validation or seed sweep, so treat the last decimal place as noise.
- The dataset is clean. Oxford-IIIT Pet photos are well-lit, centred and mostly single-subject. Expect worse on cluttered phone snapshots.
- Grad-CAM is a coarse explanation. It shows where the network responded in the last convolutional layer, not why in any causal sense. It is a debugging aid — good for catching a model that keys on the background — not proof of reasoning.
uv run train.py # auto-detects a GPU; falls back to CPU
uv run evaluate.py # test metrics + regenerates both figurestrain.py downloads Oxford-IIIT Pet (~800 MB) into data/ on first run via
torchvision. The dataset is not redistributed here.
pyproject.toml deliberately pins CPU-only torch, because that is what the
inference container needs and it installs anywhere. PyPI's default torch wheel
on Windows is CPU-only, so torch.cuda.is_available() reports False even on a
perfectly good NVIDIA card. To train on a GPU, install a CUDA build over the top:
uv pip install torch torchvision --index-url https://download.pytorch.org/whl/cu128 --reinstall--reinstall matters. Without it uv sees the version constraint already
satisfied and silently keeps the CPU wheel.
Useful flags: --epochs, --warmup-epochs, --finetune-lr, --head-lr,
--batch-size, --size, --device, --no-amp, --workers.
Backbone — clf/model.py. efficientnet_b0 with ImageNet
weights, its 1000-way head replaced by dropout plus a 37-way linear layer.
4.05M parameters. EfficientNet-B0 is chosen over a ResNet largely for size: the
checkpoint is 15.8 MB and lives in the repository, where a ResNet-50 would
be ~100 MB and hit GitHub's file limit.
Data — clf/data.py. The dataset's official trainval/test division is preserved; the validation slice is carved out of trainval only, so no checkpoint is ever selected using test data. Validation images are served through the evaluation transform, not the training one, so val scores mean the same thing test scores do.
Augmentation — resized crops kept fairly tight (scale 0.7–1.0), horizontal flips, mild colour jitter, small rotations and light random erasing. Crops are deliberately not aggressive: several breed pairs here differ only by coat markings, and hard cropping throws that signal away.
Loss — cross-entropy with 0.1 label smoothing, which helps on a fine-grained problem where several classes genuinely overlap.
Grad-CAM — clf/gradcam.py. Forward and full-backward hooks on the last convolutional stage. The gradient of the class score with respect to that feature map gives a per-channel importance weight; the weighted, ReLU'd sum is the evidence map, upsampled over the original image.
POST /api/classify — multipart: file.
GET /api/health — model status, parameter count, class list and test metrics.
The checkpoint loads on a background thread at startup, so the page is usable
immediately and reports load state in the header.
Uploads are capped at 15 MB; images longer than 900 px on an edge are downscaled first.
| Path | Role |
|---|---|
| clf/model.py | Backbone, head, freeze/unfreeze, checkpoint loading |
| clf/data.py | Dataset, splits, augmentation |
| clf/train.py | Two-phase training loop |
| clf/metrics.py | Top-k, confusion matrix, macro-F1 |
| clf/gradcam.py | Grad-CAM implementation |
| clf/inference.py | Checkpoint loading and prediction |
| clf/render.py | Heat-map colouring and blending |
| clf/server.py | FastAPI routes |
| clf/static/ | Frontend — no build step, no JS dependencies |
| train.py / evaluate.py | CLI entry points |
Code is MIT.
Trained on the Oxford-IIIT Pet Dataset (Parkhi et al., CVPR 2012), used
under CC BY-SA 4.0. Four test-split images are bundled in clf/static/samples/
so the app works straight after cloning, and the figures in docs/ derive from
the same dataset. Full citation and per-file provenance in
docs/ATTRIBUTION.md.
{ "predictions": [ { "label": "Bengal", "probability": 0.9713 }, { "label": "Egyptian Mau", "probability": 0.0121 } ], "margin": 0.9592, // top-1 minus top-2, a simple confidence signal "original": "data:image/jpeg;base64,...", "explanation": "data:image/jpeg;base64,...", "width": 500, "height": 375 }