Some further optimizations.

extras/easypose/utils.py

@@ -65,6 +65,61 @@ def nms(boxes: np.ndarray, iou_threshold: float, conf_threshold: float):
     return np.array(result)
 
 
+def nms_optimized(boxes: np.ndarray, iou_threshold: float, conf_threshold: float):
+    """
+    Perform Non-Maximum Suppression (NMS) on bounding boxes for a single class.
+    """
+
+    # Filter out boxes with low confidence scores
+    scores = boxes[:, 4]
+    keep = scores > conf_threshold
+    boxes = boxes[keep]
+    scores = scores[keep]
+
+    if boxes.shape[0] == 0:
+        return np.empty((0, 5), dtype=boxes.dtype)
+
+    # Compute the area of the bounding boxes
+    x1 = boxes[:, 0]
+    y1 = boxes[:, 1]
+    x2 = boxes[:, 2]
+    y2 = boxes[:, 3]
+    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
+
+    # Sort the boxes by scores in descending order
+    order = scores.argsort()[::-1]
+
+    keep_indices = []
+    while order.size > 0:
+        i = order[0]
+        keep_indices.append(i)
+
+        # Compute IoU of the current box with the rest
+        xx1 = np.maximum(x1[i], x1[order[1:]])
+        yy1 = np.maximum(y1[i], y1[order[1:]])
+        xx2 = np.minimum(x2[i], x2[order[1:]])
+        yy2 = np.minimum(y2[i], y2[order[1:]])
+
+        # Compute width and height of the overlapping area
+        w = np.maximum(0.0, xx2 - xx1 + 1)
+        h = np.maximum(0.0, yy2 - yy1 + 1)
+
+        # Compute the area of the intersection
+        inter = w * h
+
+        # Compute the IoU
+        iou = inter / (areas[i] + areas[order[1:]] - inter)
+
+        # Keep boxes with IoU less than the threshold
+        inds = np.where(iou <= iou_threshold)[0]
+
+        # Update the order array
+        order = order[inds + 1]
+
+    # Return the boxes that are kept
+    return boxes[keep_indices]
+
+
 def get_heatmap_points(heatmap: np.ndarray):
     keypoints = np.zeros([1, heatmap.shape[0], 3], dtype=np.float32)
     for i in range(heatmap.shape[0]):