CP-Cluster: A new box clustering strategy to replace NMS

In CVPR 2022, NVIDIA proposed a new post-processing framework for object detectors to replace NMS-based methods, named CP-Cluster:

Confidence Propagation Cluster: Unleash Full Potential of Object Detectors(arxiv,[code])(https://github.com/shenyi0220/CP-Cluster)

Motivations

Compared with NMS series, CP is fully parallelizable and achieves better accuracies when paired with most mainstream detectors.

The rule-based NMS/Soft-NMS are not perfect. From authors of the paper, it suffers from below drawbacks:

1. Non-parallelism. As the step of sorting of bounding boxes by confidence scores is necessary for all NMS-based methods.
2. Previously only suppression is considered, while the enhancement for those important boxes is fully neglected.
3. Author also mentioned that the box with highest confidence score may not the best choice to get highest IOU with GT.

So CP-Cluster is designed to to be fully parallelizable and meanwhile gain better accuracy when paired with those mainstream detectors.

Methods

Generally, CP borrows the idea of message propagation from "Belief Propagation" to propagate messages among neighboring boxes(highly overlapped boxes). By this way, redundant boxes are suppressed and those important boxes are enhanced simultaneously. The general pipeline of CP is illustrated by below diagram:

Firstly, CP build graph models for all candidate boxes outputted by a detector:

After generating graphs, CP update each box with positive messages and negative messages from neighboring boxes:

Positive messages and negative messages are generated by below 2 formulas separately:

I ignore more details for those formulas and directly started posting latest experimental results, as those improved numbers are the most attractive part for CP.

Experiment Results

Yolov5 v6.1
From the code repo of CP, it shows stable improvements on Yolov5 models. When paired with CP, the strongest model(Yolov5X_6) could achieve 56.2(mAP on COCO val) with TTA

YoloX
When applied to YoloX, CP could achieve similar accuracy improvements

More Detectors from MMDetection

CenterNet
The authors showed off CP with CenterNet, as CenterNet is among the first wave of detectors with no need of post-processing methods.