How do you go from a classifier to finding objects in images?

"Is there a cat in this picture?" and "where is every cat in this picture?" feel like different problems. In code, they are surprisingly close. The first network ends by averaging all it found across the image into one answer. The second keeps the answers separate, place by place.

This project generates its own images: circles, squares and triangles on a noisy gray background, 64 by 64 pixels. Because a script draws them, you get perfect labels, unlimited data and no rights to worry about, and the training takes seconds. The script is 05_classify_then_detect.py; it needs PyTorch and NumPy, and uses a GPU if you have one. The results below ran on an Apple M5 Max's built-in GPU with PyTorch 2.15. Other machines were not tested.

Part 1: the classifier

The body is six small convolution layers, filters that slide across the image, with three halvings of the image size. A 64 by 64 image becomes 8 by 8 cells, each described by 64 numbers.

import torch.nn as nn

def backbone():
    layers, ch = [], 1
    for out in (16, 32, 64):        # three halvings: 64 -> 8 cells
        layers += [nn.Conv2d(ch, out, 3, padding=1), nn.BatchNorm2d(out), nn.ReLU(),
                   nn.Conv2d(out, out, 3, padding=1), nn.BatchNorm2d(out), nn.ReLU(),
                   nn.MaxPool2d(2)]
        ch = out
    return nn.Sequential(*layers)

class Classifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.body, self.head = backbone(), nn.Linear(64, 3)
    def forward(self, x):
        return self.head(self.body(x).mean((2, 3)))   # average away the "where"

After 600 steps and 4.7 seconds, with one shape in each image, this model of 72,435 parameters got 1,000 of 1,000 new test images right. Then the script showed it 200 images holding three shapes of at least two kinds. It kept answering with an average top probability of 94%, one label per image. It is confident, and the question it answers has stopped fitting the picture. That .mean((2, 3)) is the reason: it throws location away.

Part 2: keep the grid

A detector, in the spirit of YOLO, replaces that average and treats the job as predicting numbers for fixed regions of the image, all in one pass. Here, each of the 64 cells predicts seven values.

class Detector(nn.Module):
    def __init__(self):
        super().__init__()
        self.body = backbone()
        self.head = nn.Conv2d(64, 1 + 3 + 3, 1)
        # for each cell: object here?, x offset, y offset, size, then 3 class scores
    def forward(self, x):
        return self.head(self.body(x))

Only 260 parameters separate the two models. What changes is the target and the loss, the score training tries to lower. The cell containing a shape's centre should say "yes" and give the box and class. All the others should say "no". Most cells are empty, so without a higher weight for "yes", the model can score well by predicting nothing at all.

Afterwards, a step called non-maximum suppression keeps the strongest of any overlapping boxes.

How do you score boxes?

A box is right if it overlaps the true one enough. IoU, intersection over union, is the overlap divided by the total area covered; 1 is perfect, and 0.5 is a common pass mark. Average precision ranks all predictions by confidence and asks how precision holds up as you try to find more of the objects. mAP is its mean across classes.

The test set is 500 fresh images. Each model trained for the same 2,500 steps, 15 to 18 seconds, on a fixed pile of images:

The table holds a lesson the classifier could not show. With 100 images the model looks acceptable at the loose test, 0.888, and poor at the strict one, 0.406. Finding things is easier to learn than placing them precisely, and placement needs more data. All four models are perfect on what they trained on, so training scores say nothing.

Then the script tested the 8,000-image model on narrower and more crowded scenes:

On the CPU, one image took 0.65 milliseconds. This grid has a built-in limit: a cell holds only one object, so two objects centred in the same cell cannot both be found. The training images keep shapes apart, which is part of why these scores are so high.

What changes with real photos?

Almost everything except the structure. These scores are high because the task is easy: flat shapes, plain background, no lighting, no perspective, and test images from the very process that made the training images. They show that the method works and how to measure it. They say nothing about how a detector would do on your photos, and nothing in this project was tested on real ones.

With real images, labels become the main cost. Someone must draw every box, and must agree with colleagues on what counts as an object. Use a pretrained body: the first layers of a network trained on millions of photos already detect edges and textures, so your data trains mainly the head. Test on photos from other days, cameras and sites, since splitting images from one shoot into train and test is the leakage of project four in another form. People in images add privacy duties, so plan for consent or blur them.

Before training anything, check if a plain classifier is enough. If you only need to know "is the shelf empty?", you need no boxes. It takes far fewer labels, since one word per image is much quicker to record than a box around every object.