Bounding Box Regression with Uncertainty for Accurate Object Detection (1809.08545v3)

Abstract: Large-scale object detection datasets (e.g., MS-COCO) try to define the ground truth bounding boxes as clear as possible. However, we observe that ambiguities are still introduced when labeling the bounding boxes. In this paper, we propose a novel bounding box regression loss for learning bounding box transformation and localization variance together. Our loss greatly improves the localization accuracies of various architectures with nearly no additional computation. The learned localization variance allows us to merge neighboring bounding boxes during non-maximum suppression (NMS), which further improves the localization performance. On MS-COCO, we boost the Average Precision (AP) of VGG-16 Faster R-CNN from 23.6% to 29.1%. More importantly, for ResNet-50-FPN Mask R-CNN, our method improves the AP and AP90 by 1.8% and 6.2% respectively, which significantly outperforms previous state-of-the-art bounding box refinement methods. Our code and models are available at: github.com/yihui-he/KL-Loss

Bounding Box Regression with Uncertainty for Accurate Object Detection

The paper presents an innovative approach to object detection by introducing a novel bounding box regression loss function, KL Loss, designed to address the ambiguities inherent in ground-truth bounding boxes. This method enhances object localization accuracy by integrating the learning of bounding box transformation and localization variance, utilizing a probabilistic framework.

Key Contributions

The authors pinpoint several shortcomings in traditional bounding box regression, primarily the inability to account for ambiguous ground-truth bounding boxes arising from factors such as occlusion and unclear object boundaries. The proposed KL Loss significantly improves localization accuracy without substantial computational overhead. The work shows notable improvements in performance across various architectures and datasets, such as MS-COCO and PASCAL VOC 2007.

1. KL Loss for Bounding Box Regression: The paper introduces KL Loss, which models the bounding box prediction as a Gaussian distribution and the ground-truth bounding box as a Dirac delta function. The use of KL divergence between these distributions allows for capturing uncertainties associated with ambiguous bounding boxes, resulting in better learning and prediction accuracy. This approach contrasts with the traditional smooth L1 loss that doesn't address the inherent ambiguities.
2. Variance Voting in NMS: The paper proposes a novel post-processing technique called variance voting, improving the non-maximum suppression (NMS) process. Instead of relying solely on classification scores, variance voting uses localization variances from neighboring bounding boxes to refine predictions further. This method significantly enhances the localization accuracy demonstrated by improvements in higher intersection over union (IoU) metrics like AP and AP⁹⁰.
3. Empirical Results: The method achieves a substantial boost in object localization performance. For instance, on the MS-COCO dataset, the approach increases the Average Precision (AP) of VGG-16 Faster R-CNN from 23.6% to 29.1%. Moreover, in ResNet-50-FPN Mask R-CNN, improvements include a 1.8% rise in AP and a 6.2% increase in AP⁹⁰, outperforming existing state-of-the-art refinement methods.

Implications and Future Directions

The proposed KL Loss and associated variance voting mechanism introduce a probabilistic perspective to bounding box regression, opening avenues for further research in uncertainty modeling in object detection. Given the growing deployment of AI models in safety-critical applications such as autonomous driving and robotics, where reliable localization confidence is imperative, this work provides a foundation for more robust models.

Future research could explore extending the probabilistic approach to other domains within computer vision and integrating it with advanced architectures or multi-task learning setups. Additionally, further investigation into different uncertainty quantification methods could provide deeper insights into improving model interpretability and reliability.

In conclusion, the incorporation of uncertainty modeling into bounding box regression represents a promising advancement in object detection, offering both theoretical insights and practical benefits in improving detection accuracy while maintaining computational efficiency. This work sets the stage for continued improvements in the critical area of object detection.