add nmsboxes

add nmsboxes

models/utils.py

@@ -61,118 +61,6 @@ def sigmoid(x: ndarray) -> ndarray:
     return 1. / (1. + np.exp(-x))
 
 
-def bbox_iou(boxes1: ndarray, boxes2: ndarray) -> ndarray:
-    boxes1_area = (boxes1[..., 2] - boxes1[..., 0]) * \
-                  (boxes1[..., 3] - boxes1[..., 1])
-    boxes2_area = (boxes2[..., 2] - boxes2[..., 0]) * \
-                  (boxes2[..., 3] - boxes2[..., 1])
-    left_up = np.maximum(boxes1[..., :2], boxes2[..., :2])
-    right_down = np.minimum(boxes1[..., 2:], boxes2[..., 2:])
-    inter_section = np.maximum(right_down - left_up, 0.0)
-    inter_area = inter_section[..., 0] * inter_section[..., 1]
-    union_area = boxes1_area + boxes2_area - inter_area
-    ious = np.maximum(1.0 * inter_area / union_area, np.finfo(np.float32).eps)
-
-    return ious
-
-
-def batched_nms(boxes: ndarray,
-                scores: ndarray,
-                iou_thres: float = 0.65,
-                conf_thres: float = 0.25):
-    labels = np.argmax(scores, axis=-1)
-    scores = np.max(scores, axis=-1)
-
-    cand = scores > conf_thres
-    boxes = boxes[cand]
-    scores = scores[cand]
-    labels = labels[cand]
-
-    keep_boxes = []
-    keep_scores = []
-    keep_labels = []
-
-    for cls in np.unique(labels):
-        cls_mask = labels == cls
-        cls_boxes = boxes[cls_mask]
-        cls_scores = scores[cls_mask]
-
-        while cls_boxes.shape[0] > 0:
-            max_idx = np.argmax(cls_scores)
-            max_box = cls_boxes[max_idx:max_idx + 1]
-            max_score = cls_scores[max_idx:max_idx + 1]
-            max_label = np.array([cls], dtype=np.int32)
-            keep_boxes.append(max_box)
-            keep_scores.append(max_score)
-            keep_labels.append(max_label)
-            other_boxes = np.delete(cls_boxes, max_idx, axis=0)
-            other_scores = np.delete(cls_scores, max_idx, axis=0)
-            ious = bbox_iou(max_box, other_boxes)
-            iou_mask = ious < iou_thres
-            if not iou_mask.any():
-                break
-            cls_boxes = other_boxes[iou_mask]
-            cls_scores = other_scores[iou_mask]
-
-    if len(keep_boxes) == 0:
-        keep_boxes = np.empty((0, 4), dtype=np.float32)
-        keep_scores = np.empty((0, ), dtype=np.float32)
-        keep_labels = np.empty((0, ), dtype=np.float32)
-
-    else:
-        keep_boxes = np.concatenate(keep_boxes, axis=0)
-        keep_scores = np.concatenate(keep_scores, axis=0)
-        keep_labels = np.concatenate(keep_labels, axis=0)
-
-    return keep_boxes, keep_scores, keep_labels
-
-
-def nms(boxes: ndarray,
-        scores: ndarray,
-        iou_thres: float = 0.65,
-        conf_thres: float = 0.25):
-    labels = np.argmax(scores, axis=-1)
-    scores = np.max(scores, axis=-1)
-
-    cand = scores > conf_thres
-    boxes = boxes[cand]
-    scores = scores[cand]
-    labels = labels[cand]
-
-    keep_boxes = []
-    keep_scores = []
-    keep_labels = []
-
-    idxs = scores.argsort()
-    while idxs.size > 0:
-        max_score_index = idxs[-1]
-        max_box = boxes[max_score_index:max_score_index + 1]
-        max_score = scores[max_score_index:max_score_index + 1]
-        max_label = np.array([labels[max_score_index]], dtype=np.int32)
-        keep_boxes.append(max_box)
-        keep_scores.append(max_score)
-        keep_labels.append(max_label)
-        if idxs.size == 1:
-            break
-        idxs = idxs[:-1]
-        other_boxes = boxes[idxs]
-        ious = bbox_iou(max_box, other_boxes)
-        iou_mask = ious < iou_thres
-        idxs = idxs[iou_mask]
-
-    if len(keep_boxes) == 0:
-        keep_boxes = np.empty((0, 4), dtype=np.float32)
-        keep_scores = np.empty((0, ), dtype=np.float32)
-        keep_labels = np.empty((0, ), dtype=np.float32)
-
-    else:
-        keep_boxes = np.concatenate(keep_boxes, axis=0)
-        keep_scores = np.concatenate(keep_scores, axis=0)
-        keep_labels = np.concatenate(keep_labels, axis=0)
-
-    return keep_boxes, keep_scores, keep_labels
-
-
 def path_to_list(images_path: Union[str, Path]) -> List:
     if isinstance(images_path, str):
         images_path = Path(images_path)
@@ -197,6 +85,50 @@ def crop_mask(masks: ndarray, bboxes: ndarray) -> ndarray:
     return masks * ((r >= x1) * (r < x2) * (c >= y1) * (c < y2))
 
 
+def box_iou(box1: ndarray, box2: ndarray) -> float:
+    x11, y11, x21, y21 = box1
+    x12, y12, x22, y22 = box2
+    x1 = max(x11, x12)
+    y1 = max(y11, y12)
+    x2 = min(x21, x22)
+    y2 = min(y21, y22)
+    inter_area = max(0, x2 - x1) * max(0, y2 - y1)
+    union_area = (x21 - x11) * (y21 - y11) + (x22 - x12) * (y22 - y12) - inter_area
+    return max(0, inter_area / union_area)
+
+
+def NMSBoxes(
+        boxes: ndarray,
+        scores: ndarray,
+        labels: ndarray,
+        iou_thres: float,
+        agnostic: bool = False
+):
+    num_boxes = boxes.shape[0]
+    order = np.argsort(scores)[::-1]
+    boxes = boxes[order]
+    labels = labels[order]
+
+    indices = []
+
+    for i in range(num_boxes):
+        box_a = boxes[i]
+        label_a = labels[i]
+        keep = True
+        for j in indices:
+            box_b = boxes[j]
+            label_b = labels[j]
+            if not agnostic and label_a != label_b:
+                continue
+            if box_iou(box_a, box_b) > iou_thres:
+                keep = False
+        if keep:
+            indices.append(i)
+
+    indices = np.array(indices, dtype=np.int32)
+    return order[indices]
+
+
 def det_postprocess(data: Tuple[ndarray, ndarray, ndarray, ndarray]):
     assert len(data) == 4
     num_dets, bboxes, scores, labels = (i[0] for i in data)