Added bbox and mask class agnostic nms code

python_utils/get_model.py

@@ -33,6 +33,57 @@ def load_model():
 
     return predictor
 
+def mask_nms(masks, scores, nms_threshold=0.5):
+    import supervision as sv
+    
+    polygons = []
+    for mask in masks:
+        contour = sv.mask_to_polygons(mask)
+        if len(contour) > 0:
+            polygons.append(Polygon(contour[0]))
+        else:
+            polygons.append(Polygon([]))
+    order = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)
+    keep = []
+    while order:
+        i = order.pop(0)
+        keep.append(i)
+        for j in order:
+            # Calculate the IoU between the two polygons
+            intersection = polygons[i].intersection(polygons[j]).area
+            union = polygons[i].union(polygons[j]).area
+            iou = intersection / union
+
+            # intersection = masks[i] * masks[j]
+            # union = masks[i] + masks[j]
+            # iou = intersection.sum() / union.sum()
+
+            # Remove masks with IoU greater than the threshold
+            if iou > nms_threshold:
+                order.remove(j)
+    return keep
+
+def apply_nms(prediction, cls_agnostic_nms=0.5, mask=False):
+    from torchvision.ops import nms
+    from detectron2.structures import Instances
+    
+    if mask:
+        # print("Applying mask NMS")
+        nms_indices = mask_nms(prediction["instances"].pred_masks.numpy(),
+                               prediction["instances"]._fields["scores"], cls_agnostic_nms)
+    else:
+        # print("Applying box NMS")
+        nms_indices = nms(prediction["instances"].pred_boxes.tensor,
+                          prediction["instances"].scores, cls_agnostic_nms)
+
+    pred = {"instances": Instances(image_size=prediction["instances"].image_size,
+                                   pred_boxes=prediction["instances"].pred_boxes[nms_indices],
+                                   scores=prediction["instances"].scores[nms_indices],
+                                   pred_classes=prediction["instances"].pred_classes[nms_indices]+1,
+                                   pred_masks=prediction["instances"].pred_masks[nms_indices])}
+
+    return pred
+
 if __name__ == '__main__':
     # get_set_up()
     load_model()