A complete object detection project — fine-tuning a pretrained detector on a custom dataset and evaluating it with mAP, the standard detection metric.
Problem Statement
Build an object detector that localizes and classifies objects within images from a small, custom dataset (e.g. detecting a handful of specific object types relevant to a chosen domain), evaluated using mean Average Precision (mAP).
Dataset
A small custom dataset annotated with bounding boxes (using a tool like LabelImg or Roboflow to create annotations if starting from raw images), or a subset of a public detection dataset like Pascal VOC for a first pass without needing custom annotation.
Architecture & Approach
Rather than implementing a detector like Faster R-CNN from scratch, this project fine-tunes a pretrained detection model from torchvision on the custom dataset — directly mirroring how object detection is approached in real practice, where training a detector from scratch is rarely the practical choice.
Step-by-Step Build
import torch
import torchvision
from torchvision.models.detection import fasterrcnn_resnet50_fpn_v2
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
# 1. Load a pretrained detector and replace its classification head for the new classes
num_classes = 4 # 3 object classes + background
model = fasterrcnn_resnet50_fpn_v2(weights='DEFAULT')
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
# 2. Custom Dataset returning images and target dicts with boxes/labels
class DetectionDataset(torch.utils.data.Dataset):
def __getitem__(self, idx):
image = load_image(self.image_paths[idx]) # a tensor, C x H x W
target = {
"boxes": torch.tensor(self.boxes[idx], dtype=torch.float32), # [[x1,y1,x2,y2], ...]
"labels": torch.tensor(self.labels[idx], dtype=torch.int64)
}
return image, target
def __len__(self): return len(self.image_paths)
def collate_fn(batch):
return tuple(zip(*batch)) # detection models expect a list of images/targets, not a stacked tensor
train_loader = torch.utils.data.DataLoader(
DetectionDataset(), batch_size=4, shuffle=True, collate_fn=collate_fn
)
# 3. Train -- detection models compute their own loss internally when given targets
optimizer = torch.optim.SGD(model.parameters(), lr=0.005, momentum=0.9)
model.train()
for epoch in range(10):
for images, targets in train_loader:
images = [img.to(device) for img in images]
targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
loss_dict = model(images, targets) # returns a dict of losses, not predictions
total_loss = sum(loss for loss in loss_dict.values())
optimizer.zero_grad()
total_loss.backward()
optimizer.step()
print(f"Epoch {epoch+1}: total_loss={total_loss.item():.4f}")
# 4. Inference -- model returns predictions when given images WITHOUT targets
model.eval()
with torch.no_grad():
predictions = model([test_image.to(device)])
print(predictions[0]['boxes'], predictions[0]['labels'], predictions[0]['scores'])Expected Results
With a modest custom dataset (a few hundred annotated images) and 10-20 epochs of fine-tuning, expect mAP in a moderate range that improves noticeably with more annotated data — object detection is generally more data-hungry than classification, so results will scale with dataset size more visibly than in the Image Classification project.
Key Learnings & Extensions
- Notice that PyTorch's detection models have a different calling convention than classification models — they compute loss internally when given targets during training, and return structured predictions (boxes, labels, scores) at inference time.
- Extension: Compute mAP properly using
torchmetrics's detection metrics rather than just visually inspecting predictions. - Extension: Compare Faster R-CNN's accuracy/speed tradeoff against a one-stage detector like YOLO on the same dataset, directly illustrating the two-stage vs one-stage detector distinction from the Computer Vision category.
