Fix fallback nms to output same value as original

mobile_sam/utils/torch_nms.py

@@ -3,13 +3,18 @@ from torchvision.ops.boxes import box_iou
 
 
 def nms(bboxes: torch.Tensor, scores: torch.Tensor, iou_threshold: float) -> torch.Tensor:
-    order = torch.argsort(-scores).to(bboxes.device)
-    indices = torch.arange(bboxes.shape[0]).to(bboxes.device)
-    keep = torch.ones_like(indices, dtype=torch.bool).to(bboxes.device)
-    for i in indices:
-        if keep[i]:
-            bbox = bboxes[order[i]]
-            iou = box_iou(bbox[None, ...], (bboxes[order[i + 1:]]) * keep[i + 1:][..., None])
-            overlapped = torch.nonzero(iou > iou_threshold)
-            keep[overlapped + i + 1] = 0
-    return order[keep]
+    order = torch.argsort(-scores)
+    keep = []
+
+    while order.numel() > 0:
+        i = order[0]
+        keep.append(i.item())
+
+        if order.numel() == 1:
+            break
+
+        ious = box_iou(bboxes[i].unsqueeze(0), bboxes[order[1:]])[0]
+        mask = ious <= iou_threshold
+        order = order[1:][mask]
+
+    return torch.tensor(keep, device=bboxes.device)

sam2/utils/torch_nms.py

@@ -3,13 +3,18 @@ from torchvision.ops.boxes import box_iou
 
 
 def nms(bboxes: torch.Tensor, scores: torch.Tensor, iou_threshold: float) -> torch.Tensor:
-    order = torch.argsort(-scores).to(bboxes.device)
-    indices = torch.arange(bboxes.shape[0]).to(bboxes.device)
-    keep = torch.ones_like(indices, dtype=torch.bool).to(bboxes.device)
-    for i in indices:
-        if keep[i]:
-            bbox = bboxes[order[i]]
-            iou = box_iou(bbox[None, ...], (bboxes[order[i + 1:]]) * keep[i + 1:][..., None])
-            overlapped = torch.nonzero(iou > iou_threshold)
-            keep[overlapped + i + 1] = 0
-    return order[keep]
+    order = torch.argsort(-scores)
+    keep = []
+
+    while order.numel() > 0:
+        i = order[0]
+        keep.append(i.item())
+
+        if order.numel() == 1:
+            break
+
+        ious = box_iou(bboxes[i].unsqueeze(0), bboxes[order[1:]])[0]
+        mask = ious <= iou_threshold
+        order = order[1:][mask]
+
+    return torch.tensor(keep, device=bboxes.device)

segment_anything_fb/utils/torch_nms.py

@@ -3,13 +3,18 @@ from torchvision.ops.boxes import box_iou
 
 
 def nms(bboxes: torch.Tensor, scores: torch.Tensor, iou_threshold: float) -> torch.Tensor:
-    order = torch.argsort(-scores).to(bboxes.device)
-    indices = torch.arange(bboxes.shape[0]).to(bboxes.device)
-    keep = torch.ones_like(indices, dtype=torch.bool).to(bboxes.device)
-    for i in indices:
-        if keep[i]:
-            bbox = bboxes[order[i]]
-            iou = box_iou(bbox[None, ...], (bboxes[order[i + 1:]]) * keep[i + 1:][..., None])
-            overlapped = torch.nonzero(iou > iou_threshold)
-            keep[overlapped + i + 1] = 0
-    return order[keep]
+    order = torch.argsort(-scores)
+    keep = []
+
+    while order.numel() > 0:
+        i = order[0]
+        keep.append(i.item())
+
+        if order.numel() == 1:
+            break
+
+        ious = box_iou(bboxes[i].unsqueeze(0), bboxes[order[1:]])[0]
+        mask = ious <= iou_threshold
+        order = order[1:][mask]
+
+    return torch.tensor(keep, device=bboxes.device)

segment_anything_hq/utils/torch_nms.py

@@ -3,13 +3,18 @@ from torchvision.ops.boxes import box_iou
 
 
 def nms(bboxes: torch.Tensor, scores: torch.Tensor, iou_threshold: float) -> torch.Tensor:
-    order = torch.argsort(-scores).to(bboxes.device)
-    indices = torch.arange(bboxes.shape[0]).to(bboxes.device)
-    keep = torch.ones_like(indices, dtype=torch.bool).to(bboxes.device)
-    for i in indices:
-        if keep[i]:
-            bbox = bboxes[order[i]]
-            iou = box_iou(bbox[None, ...], (bboxes[order[i + 1:]]) * keep[i + 1:][..., None])
-            overlapped = torch.nonzero(iou > iou_threshold)
-            keep[overlapped + i + 1] = 0
-    return order[keep]
+    order = torch.argsort(-scores)
+    keep = []
+
+    while order.numel() > 0:
+        i = order[0]
+        keep.append(i.item())
+
+        if order.numel() == 1:
+            break
+
+        ious = box_iou(bboxes[i].unsqueeze(0), bboxes[order[1:]])[0]
+        mask = ious <= iou_threshold
+        order = order[1:][mask]
+
+    return torch.tensor(keep, device=bboxes.device)