Add NMS version which can do confidence based thresholding #42370

🚀 Feature

Add NMS version which takes confidence threshold and prunes out low confidence scores from the entire prediction.

Motivation

Motivation is that if Pytorch NMS requires only boxes which cross confidence threshold, it means that the filtering should be done apriori. This filtering usually results in an ONNX NONZERO operator, and if using GLOW, NONZERO can't be supported, owing to the dynamic shape of the output tensor. ONNX supports NonMaxSuppression node, which does exactly the same thing. If pytorch also supports this, conversion of Pytorch NMS -> ONNX NMS becomes much easier.

Pitch

Add an additional input to NMS operator (confidence_threshold). New NMS operator takes entire model prediction as an input, and based on the confidence_threshold, filters out boxes lower than the threshold before moving to IOU.

Start by locating the existing PyTorch NMS operator and its ONNX conversion path. Define how a confidence_threshold input filters the full prediction before IoU suppression, then verify that the resulting behavior supports conversion to ONNX NonMaxSuppression without the proposed dynamic-shape filtering step.