Overview¶
This lesson trains a Faster R-CNN object detector on a marine debris dataset and evaluates its predictions. Faster R-CNN is a two-stage detector, and understanding those two stages is essential for working with it correctly.
Stage 1: Region Proposal Network (RPN). The RPN slides a small network across the feature map produced by the backbone. At each position it predicts whether an anchor box contains any object at all (objectness) and how much to shift the anchor to better fit that object. The output is a ranked list of candidate regions: “there is probably something here.”
Stage 2: ROI Head. Each candidate region is cropped from the feature map using ROI Align, resized to a fixed spatial extent, and passed through a classification head and a box regression head. The classification head assigns a class label. The box regression head refines the bounding box coordinates a second time.
The backbone is ResNet-50 with a Feature Pyramid Network (FPN). The FPN fuses feature maps from multiple depths of the ResNet, producing representations at several spatial scales simultaneously. This matters for marine debris detection: a candy wrapper and a derelict fishing net can appear at very different scales in the same image. Without multi-scale features, a fixed-stride backbone will miss objects at the scales it was not designed for.
torchvision ships a pretrained fasterrcnn_resnet50_fpn model. The lesson fine-tunes it on the Win2026TrashTank dataset by replacing the final classification head with one sized for the number of debris classes.
Learning Objectives¶
By the end of this lesson you will be able to:
Explain the Faster R-CNN two-stage detection pipeline, including the role of the RPN and the ROI head.
Configure torchvision’s pretrained detection models for a custom number of classes.
Write a PyTorch Dataset class that converts YOLO-format annotations to the format torchvision expects.
Train a detection model and monitor the four component losses with TensorBoard.
Run evaluation and interpret the raw predicted box output.
!pip install -q torch torchvision huggingface_hub pyyamlDataset: Win2026 TrashTank¶
The Win2026TrashTank dataset contains annotated images of marine debris items photographed in aquarium and tank settings. It was collected as part of a robotics competition in which teams built underwater vehicles capable of identifying and retrieving debris items. The controlled tank environment produces images with relatively consistent lighting and no water column particulates, which makes it a clean entry point for learning detection pipelines before moving to noisier open-ocean data.
Annotations are stored in YOLO text format. Each image has a corresponding .txt file in which each line represents one object: class_id cx cy w h, where cx cy is the bounding box center and w h is the box width and height, all normalized to the range [0, 1] relative to the image dimensions.
snapshot_download from huggingface_hub mirrors the entire dataset repository to a local directory in one call. It is equivalent to cloning the repo but handles large file storage and partial re-downloads gracefully. If the dataset repository is private, you need a Hugging Face access token. In Colab, store the token using userdata (from google.colab import userdata; token = userdata.get('HF_TOKEN')) rather than hard-coding it in the notebook.
import os
from huggingface_hub import snapshot_download
# In Colab, use userdata.get('HF_TOKEN').
import os
token = os.environ.get('HF_TOKEN', None)
print('Downloading Win2026TrashTank dataset...')
dataset_path = snapshot_download(
repo_id='OceanCV/Win2026TrashTank',
repo_type='dataset',
local_dir='./Dataset',
token=token
)
print(f"Dataset ready at: {dataset_path}")Parsing YOLO Annotations for Faster R-CNN¶
torchvision’s detection models expect targets in a specific format: a list of dicts, one per image, where each dict has a boxes key containing a float tensor of shape [N, 4] in [x_min, y_min, x_max, y_max] pixel coordinates, and a labels key containing a long tensor of integer class indices.
YOLO’s format is normalized [cx, cy, w, h], which means you need to convert before passing annotations to the model. The conversion is:
x_min = (cx - w/2) * image_width
y_min = (cy - h/2) * image_height
x_max = (cx + w/2) * image_width
y_max = (cy + h/2) * image_heightThe dataset class reads the image dimensions from the loaded PIL image, performs this conversion, and returns the box tensor in pixel space.
The class index offset is equally important. torchvision reserves label index 0 for the background class. If your YOLO annotations use 0-indexed class IDs starting at 0, you must add 1 to all class indices before passing them to the model. Passing a target with label 0 tells the model that a region is background, which is the opposite of what you want. The label dictionary is built from the YAML file that accompanies the dataset, which lists the class names in order.
The custom collate_fn is needed because different images in a batch have different numbers of objects. PyTorch’s default collation tries to stack tensors along a new batch dimension, which fails when tensor shapes differ. The custom function simply returns a list of image tensors and a list of target dicts, which is what torchvision’s detection training loop expects.
import torch
import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.transforms import functional as F
from torch.utils.data import Dataset, DataLoader
import yaml
import glob
from PIL import Image
import numpy as np
# Load classes
with open('./Dataset/data.yaml', 'r') as f:
data_yaml = yaml.safe_load(f)
classes = data_yaml['names']
num_classes = len(classes) + 1 # PyTorch reserves 0 for background
print('Classes:', classes, 'Total including background:', num_classes)
class YoloToPytorchDataset(Dataset):
def __init__(self, image_dir):
self.image_paths = sorted(glob.glob(os.path.join(image_dir, '*.jpg')))
def __len__(self):
return len(self.image_paths)
def __getitem__(self, idx):
img_path = self.image_paths[idx]
img = Image.open(img_path).convert("RGB")
w, h = img.size
txt_path = img_path.replace('.jpg', '.txt').replace('images', 'labels')
boxes = []
labels = []
if os.path.exists(txt_path):
with open(txt_path, 'r') as f:
for line in f:
parts = line.strip().split()
if len(parts) >= 5:
cls_id = int(parts[0]) + 1 # 1-indexed
cx, cy, bw, bh = map(float, parts[1:5])
xmin = (cx - bw / 2) * w
ymin = (cy - bh / 2) * h
xmax = (cx + bw / 2) * w
ymax = (cy + bh / 2) * h
boxes.append([xmin, ymin, xmax, ymax])
labels.append(cls_id)
if len(boxes) == 0:
boxes = torch.empty((0, 4), dtype=torch.float32)
labels = torch.empty((0,), dtype=torch.int64)
area = torch.empty((0,), dtype=torch.float32)
else:
boxes = torch.tensor(boxes, dtype=torch.float32)
labels = torch.tensor(labels, dtype=torch.int64)
area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
target = {}
target["boxes"] = boxes
target["labels"] = labels
target["image_id"] = torch.tensor([idx])
target["area"] = area
target["iscrowd"] = torch.zeros((len(labels),), dtype=torch.int64)
img_tensor = F.to_tensor(img)
return img_tensor, target
def collate_fn(batch):
return tuple(zip(*batch))
train_dataset = YoloToPytorchDataset('./Dataset/train/images')
val_dataset = YoloToPytorchDataset('./Dataset/val/images')
train_loader = DataLoader(train_dataset, batch_size=4, shuffle=True, collate_fn=collate_fn)
val_loader = DataLoader(val_dataset, batch_size=4, shuffle=False, collate_fn=collate_fn)
print("Data Loaders Ready.")Monitoring Training with TensorBoard¶
TensorBoard visualizes scalar metrics logged during training as interactive time-series plots. Launch it before starting the training loop so it captures data from the very first epoch rather than catching up partway through.
In Colab, %tensorboard --logdir ./runs/resnet_detect opens an embedded TensorBoard panel in the notebook output. In a local environment, run tensorboard --logdir ./runs/resnet_detect in a terminal and open http://localhost:6006 in a browser.
Faster R-CNN returns four losses on each forward pass, and all four are worth monitoring:
loss_classifier: cross-entropy loss from the ROI head’s class prediction. A plateau here whileloss_objectnessis still falling often indicates that proposals are being found but are not being classified correctly.loss_box_reg: smooth L1 loss on the ROI head’s bounding box regression. This measures how well the stage-2 refinement is converging.loss_objectness: binary cross-entropy on the RPN’s foreground/background prediction. This is typically the first loss to drop because “is there something here” is an easier question than “what is it.”loss_rpn_box_reg: smooth L1 loss on the RPN’s anchor adjustments. This governs how well stage 1 proposes tight candidate regions.
The total loss passed to the optimizer is the sum of these four components. Watching them separately tells you which part of the two-stage pipeline is underperforming.
from torch.utils.tensorboard import SummaryWriter
%load_ext tensorboard
%tensorboard --logdir ./runs/resnet_detect
writer = SummaryWriter('./runs/resnet_detect')Training Faster R-CNN¶
fasterrcnn_resnet50_fpn(pretrained=True) loads a model with weights pretrained on the COCO dataset. The backbone and RPN weights are already well-initialized; only the ROI head classifier needs to be replaced because the number of classes differs between COCO and the TrashTank dataset.
The replacement is done by reading model.roi_heads.box_predictor.cls_score.in_features to find the number of input channels, then constructing a new FastRCNNPredictor(in_features, num_classes + 1). The + 1 accounts for the background class at index 0. Replacing just this final head is a form of transfer learning: the backbone continues to provide strong general-purpose visual features, and only the task-specific classification layer needs to be trained from scratch.
SGD with momentum (0.9) is the standard optimizer for this architecture in the torchvision detection reference implementation. A learning rate in the range 0.005 to 0.02 works well for fine-tuning. Weight decay (1e-4) provides regularization.
The training loop differs from classification loops in one important way: the model is called with both images and targets during training, and it returns a dict of losses rather than predictions. There is no separate criterion function call. The total loss is computed by summing the dict values: sum(loss_dict.values()). This is then backpropagated in the usual way. During inference the targets argument is omitted entirely, and the model returns predictions instead.
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
print("Training on:", device)
# Load pre-trained Faster R-CNN
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
# Replace the head
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
model.to(device)
# Optimizer
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005)
# Note: Training for 1 epoch for demonstration
num_epochs = 1
for epoch in range(num_epochs):
model.train()
epoch_loss = 0
for images, targets in train_loader:
images = list(image.to(device) for image in images)
targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
loss_dict = model(images, targets)
losses = sum(loss for loss in loss_dict.values())
optimizer.zero_grad()
losses.backward()
optimizer.step()
epoch_loss += losses.item()
avg_loss = epoch_loss / len(train_loader)
writer.add_scalar('Loss/train', avg_loss, epoch)
print(f"Epoch {epoch+1} | Loss: {avg_loss:.4f}")
writer.close()
print("Training Complete!")Evaluating the Model¶
Switching the model to eval() mode disables training-specific behaviors: batch normalization layers use their running statistics rather than batch statistics, and dropout layers are deactivated. Both changes are necessary to get stable, reproducible predictions during inference.
Wrapping the evaluation loop in torch.no_grad() prevents gradient computation, which saves memory and speeds up inference. This is important during evaluation because the forward pass stores intermediate activations for the backward pass by default, and those tensors are not needed when you are only predicting.
The output predictions[0]['boxes'] is a tensor of shape [N, 4], where N is the number of detections above the model’s internal confidence threshold. In addition to boxes, the prediction dict contains labels (integer class indices) and scores (confidence values between 0 and 1). The model applies Non-Maximum Suppression internally before returning these outputs, so duplicated boxes for the same object are already filtered.
For a complete evaluation, you would compute mean Average Precision (mAP) by comparing predicted boxes to ground-truth boxes across IoU thresholds. torchvision provides torchvision.ops.box_iou for computing intersection over union between two sets of boxes. The torchmetrics library provides a MeanAveragePrecision metric class that handles the full COCO-style mAP calculation, including per-class AP and area-based breakdowns.
model.eval()
with torch.no_grad():
for images, targets in val_loader:
images = list(img.to(device) for img in images)
predictions = model(images)
print("Sample Prediction Boxes:", predictions[0]['boxes'].shape)
breakReflecting on Results¶
Consider the following questions as you review your training curves and evaluation output:
Speed vs. accuracy trade-offs. Faster R-CNN runs two forward passes per image: one through the RPN and one through the ROI head. YOLO processes each image in a single pass, predicting boxes and class probabilities simultaneously from a grid of anchor points. Given this architectural difference, when would you choose Faster R-CNN over YOLO for a marine monitoring application, and when would you choose YOLO?
Challenges specific to marine debris. Standard COCO categories (person, car, bicycle) appear consistently in natural light with clear edges and predictable aspect ratios. Marine debris in a tank or underwater may be partially occluded by sand, draped over rocks, coated in biofouling, or entangled with other objects. Which of these factors do you think most affects detection performance, and how would you address it with data or model changes?
Generalizing from tank to ocean. The Win2026TrashTank dataset was captured in a controlled tank with artificial lighting. If you wanted to deploy this model on footage from an ROV operating in open water, what domain shift problems would you expect to encounter, and what steps (additional data, augmentation, domain adaptation techniques) would you take to address them?