Copyright (c) Microsoft Corporation. All rights reserved.

Licensed under the MIT License.

Training an Object Detection Model¶

In this notebook, we give an introduction to training an object detection model using torchvision. Using a small dataset, we demonstrate how to train and evaluate a FasterRCNN object detection model. We also cover one of the most common ways to store data on a file system for this type of problem.

To learn more about how object detection work, visit our FAQ.

Initialization¶

Import all the functions we need.

This shows your machine's GPUs (if it has any) and the computing device torch/torchvision is using.

Next, set some model runtime parameters. We use the unzip_url helper function to download and unzip the data used in this example notebook.

Prepare Object Detection Dataset¶

In this notebook, we use a toy dataset called Fridge Objects, which consists of 134 images of 4 classes of beverage container {can, carton, milk bottle, water bottle} photos taken on different backgrounds. The helper function downloads and unzips data set to the ComputerVision/data directory.

Set that directory in the path variable for ease of use throughout the notebook.

You'll notice that we have two different folders inside:

• /images
• /annotations

This format for object detection is fairly common.

/data
+-- images
|   +-- image1.jpg
|   +-- image2.jpg
|   +-- ...
+-- annotations
|   +-- image1.xml
|   +-- image2.xml
|   +-- ...
+-- ...

Each image corresponds to an xml file. Each xml file contains information on where its corresponding image file is located. It also contains information about the bounding boxes and the object labels. In this example, our fridge object dataset is annotated in Pascal VOC format. This is one of the most common formats for labelling object detection datasets.

<!-- Example Pascal VOC annotation -->
<annotation>
    <folder>images</folder>
    <filename>1.jpg</filename>
    <path>../images/1.jpg</path>
    <source>
        <database>Unknown</database>
    </source>
    <size>
        <width>499</width>
        <height>666</height>
        <depth>3</depth>
    </size>
    <segmented>0</segmented>

    <object>
        <name>carton</name>
        <pose>Unspecified</pose>
        <truncated>0</truncated>
        <difficult>0</difficult>
        <bndbox>
            <xmin>100</xmin>
            <ymin>173</ymin>
            <xmax>233</xmax>
            <ymax>521</ymax>
        </bndbox>
    </object>
</annotation>

You'll notice that inside the annotation xml file, we can see which image the file references <path>, the number of <objects> in the image, that the image is of (<name>) and the bounding box of that object (<bndbox>).

COCO format¶

A second popular annotation standard uses the COCO format introduced by the COCO challenge. Annotations in this format need to be converted first to Pascal VOC syntax in order to run this notebook. The function coco2voc does exactly that. See the notebook 04_coco_accuracy_vs_speed.ipynb for a code example.

Load Images¶

To load the data, we need to create a Dataset object class that Torchvision knows how to use. In short, we'll need to create a class and implement the __getitem__ method. More information here: https://pytorch.org/tutorials/intermediate/torchvision_tutorial.html#defining-the-dataset

To make it more convinient, we've created a DetectionDataset class that knows how to extract annotation information from the Pascal VOC format and meet the requirements of the Torchvision dataset object class.

Inspect our Datasets/DataLoaders to make sure the train/test split looks right.

And some more detailed information, ie. the counts of ground truth boxes per class, as well as the distribution of absolute and relative widths/heights of the annotations.

By default, we apply some image transformations to the training set. Lets see how the transformations look:

Lets visualize the annotations to make sure they look correct.

Finetune a Pretrained Model¶

For the DetectionLearner, we use Faster R-CNN as the default model, and Stochastic Gradient Descent as our default optimizer.

Our Faster R-CNN model is pretrained on COCO, a large-scale object detection, segmentation, and captioning dataset that contains over 200K labeled images with over 80 label cateogories.

When we initialize the DetectionLearner, we can pass in the dataset. By default, the object will set the model to torchvision's Faster R-CNN. Alternatively, we can pass in the model we want to use - we would customize the model using the get_pretrained_fasterrcnn function and specify the parameters we want.

Computation of model accuracy requires a GPU; training does not, although without GPU training is too slow for any practical purpose. To also support CPU-only machines, we skip computation steps which would otherwise throw an error.

Fine-tune the model using our training data loader (train_dl) and evaluate our results using our testing data loader (test_dl). Here we pass in the learning rate. If not set, the learning rate will automatically be set at 0.005.

Visualize the loss and average precision (ap) over time.

Evaluate¶

We can simply run the evaluate() method on our detector to evaluate the results.

Below, you'll notice the notation IoU=0.5:0.95 in the results above. This means: the average mAP over different IoU thresholds, from 0.5 to 0.95, at step size of 0.05 (0.5, 0.55, 0.6, 0.65, 0.7, 0.75, 0.8, 0.85, 0.9, 0.95). For more information on IoU and mAP, visit our FAQ.

Plot precision-recall curves. We'll use the plot_pr_curves function to plot the PR of each IoU threshold, and also the average over the IoU thresholds. If we look at the graph on the left, we can see that the lower the IOU threshold, the higher better the results are for the PR curve. This makes sense as the model is more lenient. This is why it is also a good idea to take a look at the mean over the IOU thresholds.

Precision-Recall curves can be hard to interpret, and hence we also provide counts how many (left) images have at least one wrong detection / missed ground truth; and (right) the total number of wrong detections / missed ground truths summed over all images. These counts are shown for different detection score thresholds, ie. effectly discarding detections with confidence score lower than a given thresholds.

As can be seen, the model performs well on the test dataset for thresholds between 0.5 and 0.9, but especially at a threshold of 0.5 over-fires on the negative images (images which do not contain any of the objects) since similar images were not included in the training set. This information is especially valuable when deciding on a precision-recall operating point. To avoid over-firing, one can either add more images to the training set directly or via hard-negative mining as explained in the 12_hard_negative_sampling.ipynb notebook.

Predict Batches¶

We can predict a group of images using the predict_dl function. We pass it a dataloader that references the images we want to predict and it will perform the predictions by the batch size specified when constructing the dataloader. In this case, we'll just use the test DataLoader (test_dl) that we created above. This dataloader uses the batch size set by the parameter BATCH_SIZE.

The predict_dl function returns a list of AnnotationBboxes. Lets print out 5 of the detection values.

We can also verify that the length of the returned detections roughly equals the number of images in the test dataset.

Using a Generator Instead

Alternatively, we can use the predict_batch function if we'd rather get a generator object instead of all the predictions back in one go. This may be preferable if the dataloader you are passing in contains many images and you don't want overflow your machine's memory.

detector_generator = detector.predict_batch(<dataloader>)
one_batch_of_predictions = next(detector_generator)

We now have our predictions on the different images from our test set. Since the data is from the test set, we also know what the ground truth bounding boxes for those images are too. Lets plot the ground truth vs predictions for the images.

To plot the predictions, we'll create a simple function (get_im_and_bboxes()) to create a generator that returns the data needed to plot the ground truths vs the detections. This generator returns a new image path (im_path), ground truth bounding boxes (gt_bboxes), and detection bounding boxes (det_bboxes) each time it is iterated on.

Since these are the parameters that the plot_detection_vs_ground_truth() function takes, we can pass the entire iterator object to the plot_grid() function. For each plot it produces, it will use a new iteration from the generator.

The red boxes represents the ground truth while the green represents the predected bounding box.

Predict Single Image¶

We can use the detector to predict new images. Start by getting a new image. In this case, we'll reuse a random image in our dataset and pretend it is a new image.

We can use the inference function, which accepts either an image or a path to an image as input. Note that we set threshold=0.5 to suppress detections with low confidence.

Now that we have our predictions, lets plot the predictions on our image.

Save the model¶

The final step is to save the model to disk. We'll do this since the deployment notebook will deploy this model.

Conclusion¶

Using the concepts introduced in this notebook, you can bring your own dataset and train an object detector to find objects of interest in a given image.