Update python_utils/get_model.py

python_utils/get_model.py

@@ -35,29 +35,17 @@ def load_model():
 
 def mask_nms(masks, scores, nms_threshold=0.5):
     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 = []
     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()
+            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: