How to Do Image Recognition With CNNs on the COCO Dataset — a Practical, Step-By-Step Guide

# coco_multilabel_train.py
import os
import numpy as np
from PIL import Image
from tqdm import tqdm

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

from pycocotools.coco import COCO
from sklearn.metrics import average_precision_score

# -----------------------
# Config (======edit these======)
COCO_ROOT = "/path/to/coco/val2017"  # path to images dir (val2017)
ANN_FILE  = "/path/to/coco/annotations/instances_val2017.json"
BATCH_SIZE = 16
NUM_EPOCHS = 5
LR = 1e-4
MAX_SAMPLES = 2000   # set smaller for quick tests; or None to use all
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# -----------------------

class CocoMultiLabelDataset(Dataset):
    """Builds a multi-label dataset from COCO 'instances' annotations."""
    def __init__(self, root, ann_file, transforms=None, max_samples=None):
        self.root = root
        self.coco = COCO(ann_file)
        self.img_ids = self.coco.getImgIds()
        if max_samples:
            self.img_ids = self.img_ids[:max_samples]
        # get list of category ids and names in a deterministic order
        self.cat_ids = self.coco.getCatIds()
        self.cat_id_to_idx = {cid: i for i, cid in enumerate(self.cat_ids)}
        self.idx_to_cat = [self.coco.loadCats([cid])[0]['name'] for cid in self.cat_ids]
        self.num_classes = len(self.cat_ids)
        self.transforms = transforms

    def __len__(self):
        return len(self.img_ids)

    def __getitem__(self, idx):
        img_id = self.img_ids[idx]
        img_info = self.coco.loadImgs([img_id])[0]
        img_path = os.path.join(self.root, img_info['file_name'])
        image = Image.open(img_path).convert("RGB")
        ann_ids = self.coco.getAnnIds(imgIds=[img_id], iscrowd=None)
        anns = self.coco.loadAnns(ann_ids)

        label = np.zeros(self.num_classes, dtype=np.float32)
        for ann in anns:
            cid = ann['category_id']
            if cid in self.cat_id_to_idx:
                label[self.cat_id_to_idx[cid]] = 1.0

        if self.transforms:
            image = self.transforms(image)

        return image, torch.from_numpy(label)

def get_transforms(train=True):
    if train:
        return transforms.Compose([
            transforms.Resize(256),
            transforms.RandomResizedCrop(224),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225])
        ])
    else:
        return transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225])
        ])

def build_model(num_classes):
    model = models.resnet50(pretrained=True)
    in_features = model.fc.in_features
    model.fc = nn.Linear(in_features, num_classes)
    return model

def train_one_epoch(model, loader, optimizer, criterion, device):
    model.train()
    running_loss = 0.0
    for imgs, targets in tqdm(loader, desc="train"):
        imgs = imgs.to(device)
        targets = targets.to(device)
        outputs = model(imgs)
        loss = criterion(outputs, targets)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        running_loss += loss.item() * imgs.size(0)
    return running_loss / len(loader.dataset)

def evaluate(model, loader, device):
    model.eval()
    all_targets = []
    all_scores = []
    with torch.no_grad():
        for imgs, targets in tqdm(loader, desc="eval"):
            imgs = imgs.to(device)
            outputs = model(imgs)
            scores = torch.sigmoid(outputs).cpu().numpy()
            all_scores.append(scores)
            all_targets.append(targets.numpy())
    y_true = np.vstack(all_targets)
    y_scores = np.vstack(all_scores)
    aps = []
    for i in range(y_true.shape[1]):
        if np.sum(y_true[:, i]) == 0:
            aps.append(np.nan)  # no positives for this class in val set
        else:
            aps.append(average_precision_score(y_true[:, i], y_scores[:, i]))
    mean_ap = np.nanmean([a for a in aps if not np.isnan(a)])
    return mean_ap, aps

def predict_image(model, img_path, transform, idx_to_cat, device, topk=5):
    model.eval()
    img = Image.open(img_path).convert("RGB")
    x = transform(img).unsqueeze(0).to(device)
    with torch.no_grad():
        out = model(x)
        probs = torch.sigmoid(out).cpu().numpy()[0]
    topk_idx = probs.argsort()[-topk:][::-1]
    return [(idx_to_cat[i], float(probs[i])) for i in topk_idx]

def main():
    # builds datasets
    train_t = get_transforms(train=False)   # using val2017; no augmentation for simplicity
    val_t = get_transforms(train=False)
    dataset = CocoMultiLabelDataset(COCO_ROOT, ANN_FILE, transforms=train_t, max_samples=MAX_SAMPLES)
    # Split: simple 80/20 split for demo
    n = len(dataset)
    split = int(0.8 * n)
    train_subset, val_subset = torch.utils.data.random_split(dataset, [split, n - split])
    train_loader = DataLoader(train_subset, batch_size=BATCH_SIZE, shuffle=True, num_workers=2)
    val_loader = DataLoader(val_subset, batch_size=BATCH_SIZE, shuffle=False, num_workers=2)

    model = build_model(dataset.num_classes).to(DEVICE)
    criterion = nn.BCEWithLogitsLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=LR)

    for epoch in range(1, NUM_EPOCHS + 1):
        loss = train_one_epoch(model, train_loader, optimizer, criterion, DEVICE)
        mean_ap, aps = evaluate(model, val_loader, DEVICE)
        print(f"Epoch {epoch}: train_loss={loss:.4f}, val_mAP={mean_ap:.4f}")

    # example inference
    sample_img = os.path.join(COCO_ROOT, dataset.coco.loadImgs([dataset.img_ids[0]])[0]['file_name'])
    preds = predict_image(model, sample_img, val_t, dataset.idx_to_cat, DEVICE, topk=5)
    print("Top predictions for sample image:", preds)

if __name__ == "__main__":
    main()