Box Prediction Rebalancing for Training Single-Stage Object Detectors with Partially Labeled Data

Abstract

Partial labeling schemes, in which annotators may label some instances of classes of interest and not label other instances, can significantly reduce annotation budgets and enable machine learning algorithms that might otherwise be impossible. However, these schemes introduce noise that makes training machine learning models difficult. The Dataset for Underwater Substrate and Invertebrate Analysis (DUSIA) uses a partial labeling scheme for its training set, which consists of thousands of partially labeled video frames. To combat the challenge of training on partially labeled data, we propose Box Prediction Rebalancing for single-stage object detectors and test our method on YOLOv5, a state-of-the-art single-stage detector. We rebalance the percentage of positive and negative detections included in the loss computation of the end-to-end model, improving our model's performance and generalizability.