Train your own object detector with Faster-RCNN & PyTorch
By: Johannes Schmidt
Tags: Faster R Cnn, Object Detection, Pytorch, Python, Deep Learning

After working with CNNs for the purpose of 2D/3D image segmentation and writing a beginner’s guide about it, I decided to try another important field in Computer Vision (CV) — object detection. There are several popular architectures like RetinaNet, YOLO, SDD and even powerful libraries like detectron2 that make object detection incredibly easy. In this tutorial, however, I want to share with you my approach on how to create a custom dataset and use it to train an object detector with PyTorch and the Faster-RCNN architecture. I will show you how images that were downloaded from the internet can be used to generate annotations (bounding boxes) with the help of the multi-dimensional image viewer napari. The provided code is specifically written for Faster-RCNN models, but parts might work with other model architectures (e.g. YOLO) because general principles apply to all common object detection models that are based on anchor/default boxes. Due to transfer learning, you will see that training an object detector sometimes requires very few images! You can find the code and a jupyter notebook on my github repo.

For this tutorial, I am going to train a human head detector. I can imagine that this is a common task for phone camera applications: detecting human faces or heads within an image. If you want to train your own object detector, e.g. for racoon detection, car detection or whatever comes into your mind, you’re at the right place. So please go ahead. It might be useful for you.

For training and experiment management, I will use PyTorch Lightning and neptune. If you’re not familiar with these packages, do not worry, you’ll be able to implement your own training logic and choose your own experiment tracker. Here’s the table of content:

Getting images Annotating Dataset building Faster R-CNN in PyTorch Training Inference

Getting images

In order to train an object detector with a deep neural network like Faster-RCNN we require a dataset. For this, I downloaded 20 images (selfies) from the internet. You can do this manually or use web scraping techniques. All images are .jpg or .png rgb or rgba files.

Here is the full dataset:

Training, validation & test data. Image by author

Let’s assume you have downloaded your images into /heads/input . Before adding bounding boxes to our input images, we should first rename our files so that they all follow the same pattern. An example on how to rename them is shown below. This simply renames all files within a directory to something like 000.jpg, 001.png etc. You can find the function get_filenames_of_path() in the utils.py script.

Annotating

There are plenty of web tools that can be used to create bounding boxes for a custom dataset. These tools usually store the information in a or several specific files, e.g. .json or .xml files. But you could also save your annotations as python dicts if you don’t want to learn another file format. Pytorch’s Faster-RCNN implementation requires the annotations (the target in network training) to be a dict with a boxes and a labels key anyway. The boxes and labels should be torch.tensors where boxes are supposed to be in xyx2y2 format (or xyxy format as stated in their docs) and labels are integer encoded, starting at 1 (as the background is assigned 0). The easiest form to save a dict as a file is using the pickle module. Fortunately, the torch package integrates some functionality of pickle, e.g. it allows us to save a file like a dict with torch.save() and load it with torch.load() . This means that we can store the annotations that we create in a pickled file. If this annotation file happens to have the same name as the image, mapping the image to its annotation file becomes really easy and creating a dataset for neural network training as well. If you already have a labeled dataset at hand, you can skip this section.

As I recently discovered napari, a multi-dimensional image viewer for python, I decided to use it to generate the labels/annotations for my dataset. Please do not expect a full fledged, perfectly working code for creating bounding boxes. This is just me making myself familiar with napari and using for my needs. If you prefer an out-of-the-box solution, I recommend taking a look at myvision.ai.

Let’s take a look at how to generate annotation files for our heads dataset with napari. For this, I heavily made use of napari’s shapes layer and created a specific Annotator class that makes annotating much easier. You can run the code within a jupyter notebook or an IPython kernel. No need to run the magic command %gui qt , as this is automatically called before starting the qt application.

This will open the napari qt-application that shows one image at a time. You can navigate through your list of images by pressing ’n’ to get the next or ‘b’ to get the previous image (custom key-bindings).

Image by author

Now, if you would like to add a label with bounding boxes for the current shown image, just enter the following into your IPython console or jupyter notebook session.

annotator.add_class(label='head', color='red')

You just need to specify the label you want and the color. Now you can start using napari’s functionality to draw bounding boxes.

Note: Don’t worry if you accidentally click ’n’ or ‘b’ on your keyboard. The created bounding boxes are saved automatically. It also doesn’t matter if you delete the image layer, as the image is read from disk every time you display the image (e.g. by clicking ’n’ or ‘b’). However, if you delete the shape layer for a label, this information is lost for this image.

Image by author

We can create as many classes as we want. For each new class, a new shape layer is created, which means we can hide specific labels if the image is cluttered with bounding boxes. We basically can do whatever we want with the bounding boxes, e.g. changing it’s color or width etc.

annotator.add_class(label='eye', color='blue')

Image by author

If you want to export the annotations for this image, you can write the following:

annotator.export(pathlib.Path('.../some_directory')))

You could also specify a name for the annotation file. When no name is given, the image’s name is taken and the .pt extension appended. I recommend this approach as this makes resuming a labeling session with the Annotator possible. Here’s an example, where annotation_ids is a list of pathlib.Path objects, similar to image_files.

annotator = Annotator(image_ids=image_files,

annotation_ids=annotation_files)

If the annotation files are in the right format and have the same name as the image itself, these annotations will be used. You can for example start labeling a bigger dataset, export some of the annotations that you managed to create in a certain time and resume labeling at a later time.

Let’s continue to create bounding boxes for every image we have. For this tutorial, we’ll stick to our heads bounding boxes and delete the eye layer that I showed above. Once you’re satisfied with the result, you can export all annotations in one go with:

annotator.export_all(pathlib.Path('.../Heads/target'))

For this project we now have two directories, something like /heads/input and /heads/target . In /heads/input , we find all images that we downloaded and in the /heads/target directory the corresponding annotations with bounding boxes and labels that we just generated. 20 images and 20 annotation files (pickled python dicts) in total. Let’s quickly take a look at the annotation files.

This gives us the following:

keys:

dict_keys(['labels', 'boxes']) labels:

array(['head', 'head', 'head', 'head', 'head', 'head'], dtype='<U4') boxes:

[array([ 14.32894795, 217.18092301, 277.02631195, 531.98354928]),

array([199.95394483, 81.49013583, 396.43420467, 287.74013235]),

array([386.66446799, 2.24671611, 588.57235932, 247.57565934]),

array([306.33552198, 251.91776453, 510.41446591, 521.12828631]),

array([525.61183407, 266.0296064 , 741.63156727, 554.77960153]),

array([723.17762021, 116.22697735, 925.08551155, 432.11512991])]

Looks about right. The labels and boxes are both stored in numpy.ndarrays in xyxy format. Now we can build a proper dataset for network training.

Dataset building

Here’s how to build your own dataset that you can use to feed the network with batches of data. The approach is similar to my previous tutorial: 2D/3D semantic segmentation with the UNet.

Let’s take a look at the dataset class ObjectDetectionDataSet:

Builds a dataset with images and their respective targets. A target is expected to be a pickled file of a dict and should contain at least a ‘boxes’ and a ‘labels’ key. inputs and targets are expected to be a list of pathlib.Path objects. In case your labels are strings, you can use mapping (a dict) to int-encode them. Returns a dict with the following keys: ‘x’, ‘x_name’, ‘y’, ‘y_name’

To better understand the arguments, here’s some more information:

transform: transformations to be applied to the data.

use_cache: Instead of reading the data from disk every time we access them, we can iterate over the dataset once in the initialization method and store the data in memory (using multiprocessing). This is quite useful for network training, where we train in epochs.

mapping: As our labels are strings, e.g. ‘head’, we should integer encode them accordingly.

convert_to_format: If your bounding boxes happen to be in a different format, e.g. xywh, you can convert them into xyxy format with convert_to_format = ‘xyxy’.

Let’s use this class to build the dataset for our head detector.

As you can see in this example I use the class ComposeDouble. This allows to stack different transformations. Clip() is used to identify the bounding boxes that are bigger than the actual image and clips them accordingly. To augment the dataset one can use the albumentation module, for which I wrote the AlbumentationWrapper class. In order to use any numpy based function on the data, one can use the FunctionWrapper class. This wrapper takes in a function and an arbitrary number of arguments to return a functools.partial. I use Double to highlight that the data comes in input-target pairs (image + annotation) as opposed to Single. Whether the input or target should be transformed, can be specified with the boolean arguments input and target. By default, only the input is used. For more information I encourage you to take a look at transformations.py. In this example, we linearly scale our image and bring it in the right dimensional order: [C, H, W].

We can now take a look at a sample from the dataset with:

sample = dataset[1]

We can see that the sample is a dict with the keys: ‘x’, ‘x_name’, ‘y’, ‘y_name’.

sample['x'].shape

-> torch.Size([3, 710, 1024])

These transformations are, however, not the only ones. The Faster R-CNN implementation by PyTorch adds some more, which I will talk about in the next section. But first, let us again visualize our dataset. This time, we can pass the dataset as an argument with the DatasetViewer class instead of passing a list of image paths.

This will open a napari application, that we can navigate with the keyboard buttons ’n’ and ‘b’ again. There is, however, only one shape layer that contains the bounding boxes of every label. We can assign a color to different labels by passing a dict to our DatasetViewer or changing the color within the napari viewer instance. The label is shown on the top left corner of every bounding box. You probably can barely see it, as the text’s color is white by default. But you can change the size and color, either accessing the napari viewer instance directly with datasetviewer.viewer or by opening a small GUI application with

datasetviewer.gui_text_properties(datasetviewer.shape_layer)

This functions takes in the shapes layer for which we would like to change the text properties. The GUI is shown on the bottom left of the viewer and was created with magicgui.