BBAug: A Package for Bounding Box Augmentation in PyTorch | by Harpal Sahota | Towards Data Science

Links:
Paper
Tensorflow Policies
BBAug
Object detection models like many neural network models work best when trained on large amounts of data. It is often the case that there is limited data available and many researchers around the world are looking into augmentations strategies to increase the amount of data available. One such piece of research was conducted by Google’s Brain Team and published in a paper called: Learning Data Augmentation Strategies for Object Detection. In this paper, the authors determine a set of augmentations called policies which perform well for object detection problems. The policies were obtained by searching for augmentations which improves general model performance.
The authors define an augmentation policy as a set of sub-policies. As the model is training, one of these sub-policies is randomly selected and used to augment the image. Within each sub-policy are the augmentations to be applied to the image one after the other. Each transformation also has two hyperparameters: probability and magnitude. The probability states how likely this augmentation will be applied and the magnitude represents the degree of the augmentation. The code snapshot below shows the policy used in the paper:
policy = [        [('TranslateX_BBox', 0.6, 4), ('Equalize', 0.8, 10)],          [('TranslateY_Only_BBoxes', 0.2, 2), ('Cutout', 0.8, 8)],       [('Sharpness', 0.0, 8), ('ShearX_BBox', 0.4, 0)],        [('ShearY_BBox', 1.0, 2), ('TranslateY_Only_BBoxes', 0.6, 6)],   [('Rotate_BBox', 0.6, 10), ('Color', 1.0, 6)],]
There are 5 sub-policies within this policy and if we take the first sub-policy it contains the TranslateX_BBox and the Equalize augmentations. The TranslateX_BBox operation translates the image on the x-axis by a magnitude 4. The magnitude does not directly translate to pixels in this instance but is scaled to a pixel value dependent on the magnitude. This augmentation also has a probability of 0.6 implying that if this augmentation is selected there is a 60% chance of the augmentation being applied. With each augmentation having an associated probability a notion of stochasticity is introduced adding a degree of randomness to training. In total, the Brain Team have come up with a total of 4 policies: v0, v1, v2 and v3. The v0 policy is shown in the paper, the other three policies contain many more sub-policies with several different transformations. Overall the augmentations fall into three categories which the authors define as:
Colour Operations: Distort colour channels, without impacting the locations of the bounding boxes
Geometric Operations: Geometrically distort the image, which correspondingly alters the location and size of the bounding boxes.
Bounding Box Operations: Only distort the pixel content contained within the bounding box annotations
So where does BBAug come into this? BBAug is a python package which implements all the policies derived by the Google Brain Team. The package is a wrapper to make use of these policies much easier. The actual augmentations are done by the excellent imgaug package.
The policy shown above is applied to an example image and shown below. Each row is a different sub-policy and each column is a different run of the said sub-policy.
As you can see, there is a degree of variation between runs of a sub-policy therefore, adding a degree of randomness to training. This is just one of the 4 policies implemented in BBAug. For a full visualisation of all 4 policies check out the GitHub page of the package. The package also comes with several useful features such as the possibility of custom policies, and bounding boxes that fall outside of the image are automatically removed or clipped if they are partially outside the image. For example, in the image below a translation augmentation is applied pushing the bounding box partially outside of the image. You can see the new bounding box has shrunk to accommodate for this.
It is also possible to create augmentations which only affect the bounding box region. In the image below a Solarize augmentation is applied only to the bounding box areas:
If you would like to know more the package comes with a custom augmentation notebook.
How easy is it to augment a single image with a random policy? It’s as easy as this:
You first need to import the policies module from the BBAug package. Within this module are the 4 different policies available to you. In the above example policy set v0 was selected because this is the one used in the paper. The module also comes with a container class which will hold our chosen policy, select a random policy and apply the randomly chosen policy to an image. Line 7 instantiates the policy container with our selected augmentation policy. The next piece of code selects a random augmentation. Finally, we apply an augmentation to the image. With just 5 lines of code, we can apply an augmentation to an image. This is rather a simple example and how can it be expanded to training a model? Luckily, the package comes with a notebook illustrating how to incorporate BBAug into a PyTorch training pipeline.
This package implements the augmentation policies derived by the Google Brain Team. Currently, all 4 policies are implemented and the package comes with accompanying notebooks to help users integrate the policies into their PyTorch training pipelines. I highly recommend checking out the package’s GitHub page for a better understanding and examples of all the policies.