pytorch based YoloV5 solution - Global Wheat Detection

A complete pytorch pipeline for training, cross-validation and inference notebooks used in Kaggle competition Global Wheat Detection (May-Aug 2020)

Table of Contents

• Brief overview of the competition images
• Notebooks description
  • [TRAIN] notebook
  • [CV] Cross Validation notebook
  • [INFERENCE] Submission notebook
• How to use
• Improvements

Brief overview of the competition images

Wheat heads were from various sources:

A few labeled images are as shown: (Blue bounding boxes)

Notebooks description

A brief content description is provided here, for detailed descriptions check the notebook comments

[TRAIN] notebook

1. Pre-Processing:

   • Handled the noisy labels (too big/small boxes etc.)
   • Stratified 5 fold split based on source

2. Augmentations:

   • Albumentations - RandomSizedCrop, HueSaturationValue, RandomBrightnessContrast, RandomRotate90, Flip, Cutout, ShiftScaleRotate
   • Mixup - https://arxiv.org/pdf/1710.09412.pdf
     2 images are mixed
     
   • Mosaic - https://arxiv.org/pdf/2004.12432.pdf
     4 images are cropped and stitched together. YoloV5 by default has a canvas where it stitches images in size multiple of 32 pixels. For batch size = 4 the canvas looks like:
     
     for batch size = 2
     

3. Configurations:

   • Default YoloV5 configuration

4. TensorBoard Analysis:

   • YoloV5 by default uses TensorBoard during training, the best model is selected using “fitness” criteria based on following parameters:
     
     
     Some of my TensorBoard training logs can be found at TensorBoard.dev

[CV] Cross Validation notebook

1. Pre-Processing:

   • Same as in [TRAIN]

2. Test Time Augmentations:

   • Flips and Rotate
     
   • Color shift
   • Scale (scale down with padding)

3. Ensemble:

   • Support for ensembling of multiple folds of the same model
   • Non-Maximum Supression (NMS) is used to ensemble final predicted boxes

4. Automated Threshold Calculations:

   • Confidence level threshold is calculated based on ground truth labels
   • Optimal Final CV score (Metric: IoU) is obtained through this

[INFERENCE] Submission notebook

1. Test Time Augmentations:

   • Same as in [CV]

2. Pseudo Labelling:

   • Multi-Round Pseudo Labelling pipeline based on https://arxiv.org/pdf/1908.02983.pdf
   • Implemented Cross Validation calculations at the end of each round to decide the best thresholds for Pseudo Labels in the next round
   • Training pipeline same as in [TRAIN]
     

3. Post-Processing and Result:

   • Final predictions made with ensembled combinations of TTA

How to use

Just change the directories according to your environment.

Google Colab deployed versions are available for
[TRAIN]
[CV]

In case of any deprecation issues/warnings in future, use the modules available in YoloV5-Mixup folder.

Improvements

Acknowledging the shortcomings is the first step for progress. Thus, listing the possible improvements that could’ve made my Model better:

• Ensemble Multi-Model/Fold predictions for Pseudo Labels, currently single model is used to make pseudo labels. Would’ve made the model more robust to noise too.
• GAN or Style Transfer could’ve been used to produce more similar labeled images from the current train images for better generalization.
• Relabeling of noisy labels using multi-folds.