Soft Non-Maximum Suppression (Soft-NMS) for Object Detection in torchvision.ops

[abdullahejazjanjua]

🚀 The feature

I propose adding the implementation of Soft-NMS to the torchvision.ops library. Soft-NMS is an improved variant of the traditional Non-Maximum Suppression (NMS) algorithm, commonly used in object detection tasks. Unlike traditional NMS, which discards all boxes that overlap based on a fixed IoU threshold, Soft-NMS decays the score of overlapping boxes instead of removing them. This approach, as demonstrated by the authors of the original paper, has shown to improve detection performance, particularly for tightly overlapping objects.

Input:

• bboxes: Tensor of shape (N, 4), representing the bounding boxes.
• scores: Tensor of shape (N,), representing the scores of the boxes.
• sigma: A parameter that controls the decay rate based on the IoU values (default: 0.5).
• iou_thres: Threshold for determining overlapping boxes (default: 0.45).
• conf_thres: Minimum score threshold for boxes to be considered (default: 0.25).

Output:

• Indices: A tensor holding a list of indices of the selected boxes after applying Soft-NMS.

Motivation, pitch

The inclusion of Soft-NMS to the torchvision.ops library aligns with PyTorch vision's mission to allow developers to focus on building and innovating and save time on implementing standard algorithms from scratch. While Soft-NMS is an important enhancement for object detection tasks, its implementation is not currently available in torchvision.ops, requiring users to implement it themselves or rely on third-party libraries.

By integrating Soft-NMS into the core library, this will allow users to improve detection accuracy, particularly in cases where objects are tightly overlapping or partially occluded, and streamline the development process.

Hence, this is my motivation for implementing this feature.

Alternatives

No response

Additional context

No response

[NicolasHug]

Thanks for the feature request @abdullahejazjanjua .

This sounds in scope, but since it woudl be non-trivial (i.e. a fair amount of code and tests), I'm trying to make sure this is in scope. The soft nms paper dates back from 2017, are you aware of recent, popular models that successfully use soft-nms instead of nms?