Create YOLOv3 using PyTorch from scratch (Part-5)

This is Part-5 of the series on building a YOLOv3 model from scratch.

Here is an overview of the series:

1. Understand the YOLO model.
2. Build the model backbone.
3. Load pre-trained weights.
4. Get the tools ready.

5. Training data preparation: this post.

   This part will write some pre-processing codes to load the COCO detection dataset, including the images and annotation labels. It is also a good chance to test out the IoU, NMS and mAP utility functions we created in the last part.

6. Train the model.

I used this bash script from the Github repo of darknet to download the 2014 COCO detection dataset.

It is generally a good practice to examine the bash script before executing it. So download the script file, have a look, put it somewhere and execute it. It will download the 2014 COCO training and validation datasets, as well labels. The entire unzipped dataset is about 23 GB in size, so make sure your disk has enough of space.

Once it is done, we should have a folder structure like this:

coco/
   images/
       train2014/
       val2014/
   labels/
       train2014/
       val2014/
   ...

I’ve omitted other files and folders. We will be focusing on the images and labels folders alone.

It might be a good idea to symlink this coco folder into the data sub-folder under the YOLOv3_pytorch project folder. For a structure of the folder, refer back to the Create the Darknet-53 model section of part-2.

Also note that we will need to map the class names to ids and back. So also put this coco.names_.txt file into the coco folder as well.

I’m going to put the relevant code in a loader.py script in the YOLOv3_pytorch project folder. You can put these anywhere you like.

Let’s test out the above shown codes. It might be a good place to also test out the codes of IoU, NMS and mAP computations, introduced in the previous part.

Create a new evaluate.py script in the YOLOv3_pytorch project folder, with the following content:

import numpy as np
import torch
from torchvision import transforms

from config import load_config
from model import Darknet53
from utils import batch_NMS, compute_mAP, draw_predictions
from loader import create_loader

if __name__=='__main__':

    #--------------------Load model config--------------------
    CONFIG_FILE = './config/yolov3.cfg'
    net_config, module_list = load_config.parse_config(CONFIG_FILE)
    config = {'net': net_config}
    config['module_list'] = module_list
    config['width'] = 416
    config['height'] = 416
    config['n_classes'] = 80
    config['max_data_size'] = 10
    config['batch_size'] = 4
    config['is_train'] = False
    config['conf_thres'] = 0.3
    config['nms_iou_thres'] = 0.5
    config['map_iou_thres'] = 0.5
    DATA_FOLDER = './data/coco'
    WEIGHT_PATH = './yolov3.weights'

    #-------------------Create model-------------------
    model = Darknet53(config)

    #-------------------Load weights-------------------
    model.load_weights(WEIGHT_PATH)

    #----------------Turn on eval model----------------
    model.eval()

    #--------------Get dataset and dataloader--------------
    dataset, dataloader = create_loader(DATA_FOLDER, config, shuffle=False)
    id2class = dataset.id2class

    #-----------------Start prediction-----------------
    preds = []
    labels = []

    for imgii, labelii in dataloader:

        with torch.no_grad():
            yii = model(imgii)

        labelii = labelii.cpu().numpy()
        yii = yii.detach().cpu().numpy()

        # filter by NMS
        yii = batch_NMS(yii, config['conf_thres'], config['nms_iou_thres'])

        # convert to fractional coordinates and sizes, to match with labels
        yii[:, [1,3]] /= config['width']
        yii[:, [2,4]] /= config['height']

        print('## yii.shape:', yii.shape)
        print('## labelii.shape:', labelii.shape)
        preds.append(yii)
        labels.append(labelii)

        # draw predictions
        if len(yii) > 0:
            # loop through images in batch
            for jj in range(len(imgii)):
                imgjj = imgii[jj]           # get image from batch
                imgjj = transforms.ToPILImage()(imgjj)   # transform image to PIL image
                yjj = yii[yii[:, 0] == jj][:, 1:]        # get predictions from batch
                yjj[:, [0, 2]] *= config['width']        # scale back to pixel units
                yjj[:, [1, 3]] *= config['height']

                fig, ax = draw_predictions(imgjj, model.width, model.height, yjj, id2class)
                #fig.show()

                #----------------- Save plot------------
                plot_save_name = 'pred_result_%d-%d' %(len(preds), jj)
                print('\n# <predict_pretrained>: Save figure to', plot_save_name)
                fig.savefig(plot_save_name, dpi=100, bbox_inches='tight')

    preds = np.vstack(preds)
    labels = np.vstack(labels)
    print(preds.shape)
    print(labels.shape)

    mAP = compute_mAP(preds, labels, 0, config['map_iou_thres'])
    print('mAP@0.5 =', mAP)

Some points to note:

• We use many building blocks created in earlier parts:

  from config import load_config
  from model import Darknet53
  from utils import batch_NMS, compute_mAP, draw_predictions
  

  So if you haven’t got these ready, please go back to previous parts of the series first.

• We call the newly created create_loader() function to get the dataset and dataloader. For test purposes, I’m reading in only 10 images, with a batch size of 4:

      config['max_data_size'] = 10
      config['batch_size'] = 4
  

• When iterating through the dataloader, imgii is a 4D tensor with shape [B, 3, 416, 416], where B is the batch size. labelii is a 2D tensor of shape [n, 6], where n is the number of labeled objects in this batch.

• There are some back-n-forth coordinate transformations.

    yii[:, [1,3]] /= config['width']
    yii[:, [2,4]] /= config['height']
  

  This is converting the model predicted bounding box coordinates to fractional units, because the labels are using fractions, and we need them to be consistent to compute the mAP.

  When creating the plot using draw_predictions(), we convert from fractions back to pixel units again:

    imgjj = transforms.ToPILImage()(imgjj)   # transform image to PIL image
    yjj = yii[yii[:, 0] == jj][:, 1:]        # get predictions from batch
    yjj[:, [0, 2]] *= config['width']        # scale back to pixel units
    yjj[:, [1, 3]] *= config['height']
  

• To compute the mAP score, we collect all predictions in the preds list, and all labels in the labels list, convert them to np.ndarray, and call our compute_mAP() function created in the previous post.

Below are some sample image outputs. Note that the NMS process has eliminated duplicate predictions. But when there are many objects relatively packed together, they are also all preserved.

The final mAP@0.5 score is 0.53. The validation sample size is too small (we only used 10 images) for this to be really meaningful, but it is comparable with the reported offical value.

In this post we created Dataset and DataLoader classes to read and load COCO 2014 detection data, then tested out the IoU, NMS and mAP computation functions.

We have got everything ready to start training a YOLOv3 model from scratch, or do some fine-tuning with pre-trained weights. We will get into the training part in the next post.