first commit

ssd/utils/nms.py

@@ -0,0 +1,33 @@
+import warnings
+import torchvision
+
+try:
+    import torch_extension
+
+    _nms = torch_extension.nms
+except ImportError:
+    if torchvision.__version__ >= '0.3.0':
+        _nms = torchvision.ops.nms
+    else:
+        from .python_nms import python_nms
+
+        _nms = python_nms
+        warnings.warn('You are using python version NMS, which is very very slow. Try compile c++ NMS '
+                      'using `cd ext & python build.py build_ext develop`')
+
+
+def boxes_nms(boxes, scores, nms_thresh, max_count=-1):
+    """ Performs non-maximum suppression, run on GPU or CPU according to
+    boxes's device.
+    Args:
+        boxes(Tensor): `xyxy` mode boxes, use absolute coordinates(or relative coordinates), shape is (n, 4)
+        scores(Tensor): scores, shape is (n, )
+        nms_thresh(float): thresh
+        max_count (int): if > 0, then only the top max_proposals are kept  after non-maximum suppression
+    Returns:
+        indices kept.
+    """
+    keep = _nms(boxes, scores, nms_thresh)
+    if max_count > 0:
+        keep = keep[:max_count]
+    return keep