Densely-Populated Traffic Detection using YOLOv5 and Non-Maximum Suppression Ensembling

Our proposed method consists of 3 main modules. First, we acquired and preprocessed the dataset. During preprocessing, we applied augmentation, resized the images into uniform shapes, and created training and testing folds. Then, four different models were trained with these different training folds. After training, we ensembled the models using Non-Maximum Suppression [16] for final inference. A complete pipeline of our proposed methodology is illustrated in Figure 1.

For this experiment, we used the “DhakaAI” [17] dataset developed under the “Dhaka AI 2020 challenge”. The dataset consists of 3000 annotated images of traffic objects. The training dataset consisted of 21 classes. The most challenging part about the dataset is it contains images of vehicles from a different point of view. There were images from the front view, back view, side view and most importantly top view of streets. We also added around 200 new images for training to increase the sample of rare class vehicles. These new images were hand-annotated using labelImg tool [18]. Most of these images were top-view nighttime images.

For generalizing a model for object detection using deep learning architecture, a prerequisite is to have enough training examples for each class so that the model can learn properly. But, after exploring the DhakaAI dataset [17], we found that it has a huge class imbalance. The number of labels for each class is shown in Table 1. Here, some of the classes have less than 50 samples in the training dataset.

To resolve this issue, we used augmentation using tools from Roboflow ¹¹1Available at: https://roboflow.com/ and Albumentations library [19] for image augmentation. Although augmentation did not provide a very good result in the case of densely populated images, it improved the result in the case of night images.

During the exploration of the dataset, we found that there was a lot of mislabelling in the DhakaAI dataset training data. We also found that two images had different labeling of class for the same car in the same frame (illustrated in Figure 2). So, we hand-annotated all 3000 images and labeled all the mislabelled objects as well as fixed labeling of wrongly labeled objects in the image.

Another challenge in the dataset is it does not have uniform image quality. Some of the images are in landscape mode while some of the images are in portrait mode. So, we resized the images to $1024\times 1024$ pixels.

For the train and validation set split, we used the k-fold Cross-Validation technique so that our model could learn from the complete dataset. While creating the fold, we tried to make sure that images from the same frame in the train split do not occur in the validation split. Count of train-validation split for each fold is given in Table 2.

Although the key priority of our work was to localize and classify the vehicular objects on a street image, we also had to look into the inference speed so that it could be implemented real time. We had to discard R-CNN, Fast R-CNN and Faster R-CNN as it could not compete with YOLO models in both performance and inference time. YOLO on the other hand, YOLO had a much less inference time with better accuracy. Among different versions of YOLO we chose YOLOv5[5] due to its simple architecture compared to R-CNN based models. Even YOLOv5 is faster and more robust than other members of YOLO family.

Even if the author of YOLOv4[20] got official approval of YOLO, the version of YOLOv5 [5] developed by Ultralytics LLC team did not get any acknowledgement from the original author of YOLO. Still YOLOv5 provides much better performance compared to other models of YOLO family[21]. YOLOv5 inherits the advantages of YOLOv4 [20] by adding SPP-NET [22] along with some enhancement techniques. YOLOv5 has become the new state of the art for object detection[23]. YOLOv5 was mainly developed to balance between real-time performance and detection accuracy.

YOLOv5s [5] , YOLOv5m [5], YOLOv5l [5], YOLOv5x [5] are the four versions of YOLOv5 where YOLOv5s being the lightest model and YOLOv5x being the heaviest model respectively. Among all these four version, there is a trade off between the detection speed and real-time performance. The key difference among these versions are the number of feature extraction modules and convolution kernel in specific location of the network.

The network consists of three networks. These are: backbone network, neck network and detect network. Backbone network is a convolutional neural network for aggregating the fine-grained images and forming image features. Neck network is responsible for combining the image features collected by backbone network and transmitting the feature map to the detect network. The detect network is responsible for detection and classification part of the model. It applies anchor boxes on the feature map from the neck network. It also contains a softmax layer which predicts the probability of the class the bounding box surrounding the object.

For image enhancement, YOLOv5 uses mosaic data augmentation to solve the low dataset problem. It applies operations like random inversion, zooming, cropping on four images and then combines them into a single image.

In traffic detection, the core priority was to improve the performance, so we chose YOLOv5x for our training model. It contains $607$ layers with $88,568,234$ trainable parameters. The model was pre-trained using Common Object in Context (COCO) dataset [24] to detect $80$ classes. For our task, we changed the final layer to detect only $21$ classes corresponding to the $21$ vehicle classes available in the DhakaAI dataset.

To ensure the robustness and accuracy of our model, we trained 4 separate models using different sets of images. Each of the models proposes multiple bounding boxes to specify candidate regions for vehicle detection. We used Non-Maximum Suppression [16] to aggregate these bounding boxes to select the ones having the most confidence. The way it works is the system takes all the bounding boxes proposed by all four models and puts them in a priority queue sorted based on the confidence of the models predicting them. It then selects the box with the highest confidence from the queue and calculates the Intersection over Union (IoU) with the rest of the boxes. If the IoU value exceeds a certain threshold for any of the remaining boxes, that box is discarded. It then removes the bounding box with the highest confidence from the queue and adds it to the selected box list. This process is repeated until there is no bounding box remaining in the priority queue. Finally, the boxes in the selected box list are returned.

During training we had to train four different models and ensembled these four models for final output. All four of our model was trained on Google Colab[25]. Google Colab provides cloud based training utility with free GPU access for limited amount of time. For each fold of our dataset, we trained a model. The first three models were trained with image resolution of $1024\times 1024$ pixels. But the fourth model was trained with image resolution of $640\times 640$ pixels. First three fold of our dataset contained all the images while the fourth fold contained only the night images. On the dataset, it was seen that the night images itself were quite distorted noisy. So. we decided to train the night images on lower resolution as it might then focus on the the larger objects of the night images. Also thus we could train our model for longer time.

All four of our models were trained with Tesla T4 GPU which comes with $16$ GB of video memory. All four of our model was trained for around $12$ hours each.

For training, we used a YOLOv5 implementation by the Ultralytics ²²2Available at: https://github.com/ultralytics/yolov5. We used Stochastic Gradient Descent as our optimizer. Image augmentation parameters used for each of the models are given in Table 4.

To evaluate our performance, we used mean Average Precision (mAP) over training epochs following . The formula for calculating mean average precision for object detection is

where $n$ is the number of classes and $AP_{k}$ is the average precision for class $k$. Average precision (AP) is a way of summarizing the precision-recall curve into a single value representing average of all precision. The formula for calculating AP is

where $n$ is the number of thresholds and $Recall(n)=0$ and $Precision(n)=1$. We used the checkpoint where the model had most [email protected]

For inference, we ensembled the weight of all four models we trained. We used Non-Maximum Suppression during the ensemble and the confidence threshold for each predicted bounding box was set to $0.3$

We used all four training model’s weights during the final inference. We ran an inference on test data - 2 provided by DhakaAI. We hand-annotated $450$ test images and executed inference. On that test, our model achieved $[email protected]$ value of $0.458$. We also conducted inference on one of our validation sets. During validation set inference, our model achieved $[email protected]$ value of $0.883$ and $[email protected]$ value of $0.677$.

We compared the result of our model with other models of YOLO family as well as Faster R-CNN model. Table 5 shows the comparison between these models. For comparison we compared our model’s performance as well as the inference time on a single image with YOLOv3, YOLOv4 and Faster R-CNN. We trained each of these models for $12$ hours on google colab in a similar environment. The table shows that, our model has achieved the most $[email protected]$. As our proposed method ensembles $4$ different models during inference, the inference time of our solution is a little bit higher compared to other models. Still the precision performance of our model outperforms the other models.

Output of our model for different scenario is illustrated in Figure 3 and 4. Our model was able to localize and detect most of the vehicular objects for a given image taken from a different view of the street. It also performed well in the case of night images. Figure 3 illustrates the performance of our model on night images. It could locate most of the vehicles as well as properly classify those vehicles in both densely populated images and in less populated images. However, as seen in Figure 3c, our model could not detect most of the vehicles in a very low-light noisy image.

In Figure 4, we illustrated our model’s performance on images taken from a different view of the streets which shows it can locate and detect the objects properly. As seen in the figure, the model also performs well in the case of occluded objects in the image.

And our model could run inference within $0.75$ second per image. So, it could also be implemented in real-time vehicular traffic detection applications.