Update python_utils/get_model.py

python_utils/get_model.py

@@ -35,34 +35,47 @@ def load_model():
     return predictor
 
 def mask_nms(masks, scores, nms_threshold=0.5):
+    """
+    Runs class agnostic NMS on masks/segmentations instead of the bounding boxes.
+    :param      masks: (list float) List of coordinates that make up the mask output from the model.
+    :param      scores: (list float) List of corresponding confidence scores given to each mask.
+    :param      nms_threshold: (float) Threshold to apply mask-based class agnostic NMS.
+    :return     masks_kept (list float): List of masks kept after applying NMS.
+    """
     import supervision as sv
+    from shapely.geometry.polygon import Polygon
     
+    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 = []
+    masks_kept = []
     while order:
         i = order.pop(0)
-        keep.append(i)
+        masks_kept.append(i)
         for j in order:
-
-            intersection = masks[i] * masks[j]
-            union = masks[i] + masks[j]
-            iou = intersection.sum() / union.sum()
+            # Calculate the IoU between the two polygons
+            intersection = polygons[i].intersection(polygons[j]).area
+            union = polygons[i].union(polygons[j]).area
+            iou = intersection / union
 
             # Remove masks with IoU greater than the threshold
             if iou > nms_threshold:
                 order.remove(j)
-    return keep
+    return masks_kept
 
 def apply_nms(prediction, mask=False, cls_agnostic_nms=0.8):
     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)