How to use BatchedNMS plugin with retinaFace Detector

I am using RetinaFaceNet for facial detection, which provides a score, a bounding box, and 5 facial landmarks for each detected face. These are represented as three separate vectors (or matrices): a score vector, a bounding box vector, and a landmark vector. I would like to apply non-maximum suppression (NMS) to filter out overlapping bounding boxes based on the scores, which is a common post-processing step in object detection.

I am considering using the BatchedNMS Plugin in TensorRT for this task, as it provides an efficient, GPU-accelerated implementation of NMS. However, I am encountering a problem. While the BatchedNMS Plugin can provide the retained bounding boxes and their scores after NMS, it does not maintain a linkage to the original indices of these bounding boxes in the input data.

This presents a challenge because I need to keep track of the facial landmarks associated with each bounding box. After applying NMS, I don’t know how to find the corresponding landmarks for the retained bounding boxes, as the original indices are lost in the process.

I am looking for suggestions on how I can maintain the linkage between bounding boxes and their corresponding landmarks through the NMS process. Specifically, how can I find the corresponding landmarks for the bounding boxes retained after NMS when using the BatchedNMS Plugin in TensorRT? Any advice or alternative strategies would be greatly appreciated.

Attached, you’ll find the ONNX model file and a screenshot of the model’s structure. The screenshot was created with the Netron tool.

mobilenet_simplified.onnx (1.7 MB)

Please refer to below links related custom plugin implementation and sample:

While IPluginV2 and IPluginV2Ext interfaces are still supported for backward compatibility with TensorRT 5.1 and 6.0.x respectively, however, we recommend that you write new plugins or refactor existing ones to target the IPluginV2DynamicExt or IPluginV2IOExt interfaces instead.


Thanks, I will check these.