Your Guide to Object Detection with Detectron2 in PyTorch

Object detection is one of the most exciting and impactful applications of deep learning. The ability to automatically localize and classify objects in images and video has powered advances in fields ranging from autonomous driving and robotics to medical imaging and retail.

While object detection was a challenging problem requiring complex, hand-engineered pipelines just a decade ago, the rise of deep learning has led to tremendous progress. Today, state-of-the-art object detectors can achieve impressive speed and accuracy, and open-source libraries have made these cutting-edge models accessible to all.

In this post, we‘ll dive into object detection using Detectron2 – Facebook AI Research‘s next-generation library for object detection and segmentation. Building on the proven architecture of the Mask R-CNN model, Detectron2 offers modular and extensible components enabling rapid exploration and optimization of models.

Whether you‘re an ML practitioner looking to apply object detection to your own projects, or a researcher eager to push the boundaries of the field, this guide will equip you with the knowledge and practical skills to get started.

We‘ll begin by setting up Detectron2 and demonstrating its core functionality by running inference with a pre-trained model. Then, we‘ll walk through training a custom detection model on a new dataset. Along the way, we‘ll highlight tips and best practices to achieve strong results. Finally, we‘ll discuss key aspects of Detectron2‘s design and point to resources to dive even deeper.

Let‘s get started!

Setting Up Detectron2

The first step is to install Detectron2 and its dependencies. Detectron2 is powered by PyTorch and optimized for GPU acceleration, so you‘ll need a CUDA-enabled GPU and the appropriate PyTorch version for your environment.

With PyTorch installed, Detectron2 itself can be installed from Facebook‘s pre-built binaries:

!pip install detectron2 -f https://dl.fbaipublicfiles.com/detectron2/wheels/cu102/torch1.7/index.html

We can then import Detectron2 and verify the installation:

import detectron2
from detectron2.utils.logger import setup_logger
setup_logger()

import numpy as np
import cv2
import matplotlib.pyplot as plt

Inference with a Pre-trained Model

Detectron2 provides a selection of pre-trained models in its model zoo that can be used out-of-the-box for inference. These models have been trained on popular datasets like COCO and optimized for various use cases and compute requirements.

To load a model, we simply specify the config file and model weights:

from detectron2.config import get_cfg
from detectron2 import model_zoo

cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/faster_rcnn_R_101_FPN_3x.yaml"))
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Detection/faster_rcnn_R_101_FPN_3x.yaml")

predictor = DefaultPredictor(cfg)

We can then apply this model to perform object detection on an input image:

image = cv2.imread("input.jpg")
output = predictor(image)

v = Visualizer(image[:,:,::-1], metadata=MetadataCatalog.get(cfg.DATASETS.TRAIN[0]))
result = v.draw_instance_predictions(output["instances"].to("cpu"))

plt.imshow(result.get_image()[:, :, ::-1])

Detectron2 makes it simple to obtain high-quality detections with just a few lines of code! The model handles localizing objects in the image and classifying them into categories defined by the training dataset (in this case, the 80 classes of COCO).

Training a Custom Model

While pre-trained models provide a great starting point, you‘ll often want to train a detector specialized for your particular use case. Detectron2 streamlines the process of training on custom datasets.

In this example, we‘ll train a model to detect a single class (balloons) in images. The balloon dataset follows the COCO format, which Detectron2 natively supports.

After downloading and extracting the balloon dataset, we register it with Detectron2:

from detectron2.data.datasets import register_coco_instances

register_coco_instances("balloon_train", {}, "path/to/balloon/train/annotation.json", "path/to/balloon/train")
register_coco_instances("balloon_val", {}, "path/to/balloon/val/annotation.json", "path/to/balloon/val")

balloon_metadata = MetadataCatalog.get("balloon_train")

We can then define our model‘s config, again inheriting from a COCO-pretrained config:

cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.DATASETS.TRAIN = ("balloon_train",)
cfg.DATASETS.TEST = ()

cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml") 
cfg.SOLVER.MAX_ITER = 1000
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1

Several key settings to note:

• We specify the train and test sets to point to our balloon dataset
• We initialize weights from a model pre-trained on COCO for faster convergence
• We define the total iterations to train for
• We set the number of classes to 1 since we are only detecting balloons

With our dataset and config ready, we can launch training:

os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = DefaultTrainer(cfg)
trainer.resume_or_load(resume=False)
trainer.train()

After training completes, we can load the final model weights and visualize its performance on the validation set:

cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.8
predictor = DefaultPredictor(cfg)

for d in random.sample(balloon_val_dicts, 5):    
    im = cv2.imread(d["file_name"])
    outputs = predictor(im)
    v = Visualizer(im[:, :, ::-1], metadata=balloon_metadata)
    v = v.draw_instance_predictions(outputs["instances"].to("cpu"))
    plt.figure(figsize=(16, 12))
    plt.imshow(v.get_image()[:, :, ::-1])

With just a small custom dataset, we are able to train a highly accurate balloon detector! This same workflow can be applied to any custom object detection task.

Tips and Tricks

A few key tips for training strong Detectron2 models:

• Start from pre-trained model weights for faster convergence, especially when training on a small dataset
• Experiment with different learning rates and learning rate schedules to optimize convergence
• Use data augmentation (random flipping, scaling, etc.) to improve robustness, available in Detectron2‘s data loader
• Evaluate performance on a validation set and use early stopping to prevent overfitting
• Test at multiple inference resolutions to trade off speed and accuracy

Under the Hood

Detectron2‘s leading performance stems from its incorporation of best practices and optimizations from the latest object detection research. A few highlights:

• Detectron2 is built on top of PyTorch and its flexible, modular APIs make it easy to swap in custom model components and loss functions
• Synchronous Batch Norm is used to ensure consistent normalization across multiple GPUs during distributed training
• Soft NMS is used to suppress duplicate detections during inference based on overlap and classification scores
• Detectron2 supports panoptic segmentation, an extension of instance segmentation that also classifies and segments background regions
• Alternate meta-architectures and backbone networks can be dropped in to unlock higher resolutions, denser predictions, and complementary feature sets

For a deeper dive into these technical details, check out the Detectron2 documentation and the original Detectron2 paper.

What‘s Next?

Detectron2 offers a flexible foundation for object detection and segmentation that both practitioners and researchers can build on. Here a few ideas for next steps:

• Experiment with custom backbones and model architectures compatible with Detectron2 – Try techniques like data distillation and self-training to get more out of small custom datasets
• Extend Detectron2 with new data augmentations, loss functions, or post-processing steps
• Apply Detectron2 to video data for spatio-temporal action localization
• Deploy your trained Detectron2 models to a mobile device using PyTorch Mobile for real-time inference

To stay up-to-date on the latest developments, follow the Detectron2 GitHub repository, and check out new research published by FAIR. With its state-of-the-art performance and vibrant community, Detectron2 is an exciting platform to build the next wave of object detection applications!