Classification-IoU Joint Label Assignment for End-to-End Object Detection

Abstract

Using prediction-aware label assignment to choose positive samples has proven to be the key to remove Non-maximum Suppression (NMS) and reach end-to-end object detection. However, existing prediction-aware label assignment methods combine the classification cost and localization cost simply by multiplication or addition. It might produce unstable assignment results since the classification cost and localization cost are used independently to assign labels, and thus degrades the detection accuracy. In this paper, we propose a classification-IoU joint label assignment that takes advantage of the correlation between classification and localization to choose positive samples. In particular, we add category and localization consistency constraints to the cost function for label assignment, and improve the stability of label assignment. Besides, a classification-IoU joint representation is applied to improve the correlation between classification and localization in the process of training. It is trained by the loss function that considers classification and location quality simultaneously. Extensive experiments demonstrate the effectiveness of our method. On COCO dataset, the end-to-end object detector with the techniques we proposed, achieves competitive performance against the state-of-the-art object detectors.