Create box_utils.py

box_utils.py

@@ -0,0 +1,115 @@
+# SPDX-License-Identifier: MIT
+
+import numpy as np
+
+def area_of(left_top, right_bottom):
+    """
+    Compute the areas of rectangles given two corners.
+    Args:
+        left_top (N, 2): left top corner.
+        right_bottom (N, 2): right bottom corner.
+    Returns:
+        area (N): return the area.
+    """
+    hw = np.clip(right_bottom - left_top, 0.0, None)
+    return hw[..., 0] * hw[..., 1]
+
+def iou_of(boxes0, boxes1, eps=1e-5):
+    """
+    Return intersection-over-union (Jaccard index) of boxes.
+    Args:
+        boxes0 (N, 4): ground truth boxes.
+        boxes1 (N or 1, 4): predicted boxes.
+        eps: a small number to avoid 0 as denominator.
+    Returns:
+        iou (N): IoU values.
+    """
+    overlap_left_top = np.maximum(boxes0[..., :2], boxes1[..., :2])
+    overlap_right_bottom = np.minimum(boxes0[..., 2:], boxes1[..., 2:])
+
+    overlap_area = area_of(overlap_left_top, overlap_right_bottom)
+    area0 = area_of(boxes0[..., :2], boxes0[..., 2:])
+    area1 = area_of(boxes1[..., :2], boxes1[..., 2:])
+    return overlap_area / (area0 + area1 - overlap_area + eps)
+
+def hard_nms(box_scores, iou_threshold, top_k=-1, candidate_size=200):
+    """
+    Perform hard non-maximum-supression to filter out boxes with iou greater
+    than threshold
+    Args:
+        box_scores (N, 5): boxes in corner-form and probabilities.
+        iou_threshold: intersection over union threshold.
+        top_k: keep top_k results. If k <= 0, keep all the results.
+        candidate_size: only consider the candidates with the highest scores.
+    Returns:
+        picked: a list of indexes of the kept boxes
+    """
+    scores = box_scores[:, -1]
+    boxes = box_scores[:, :-1]
+    picked = []
+    indexes = np.argsort(scores)
+    indexes = indexes[-candidate_size:]
+    while len(indexes) > 0:
+        current = indexes[-1]
+        picked.append(current)
+        if 0 < top_k == len(picked) or len(indexes) == 1:
+            break
+        current_box = boxes[current, :]
+        indexes = indexes[:-1]
+        rest_boxes = boxes[indexes, :]
+        iou = iou_of(
+            rest_boxes,
+            np.expand_dims(current_box, axis=0),
+        )
+        indexes = indexes[iou <= iou_threshold]
+
+    return box_scores[picked, :]
+
+def predict(width, height, confidences, boxes, prob_threshold, iou_threshold=0.5, top_k=-1):
+    """
+    Select boxes that contain human faces
+    Args:
+        width: original image width
+        height: original image height
+        confidences (N, 2): confidence array
+        boxes (N, 4): boxes array in corner-form
+        iou_threshold: intersection over union threshold.
+        top_k: keep top_k results. If k <= 0, keep all the results.
+    Returns:
+        boxes (k, 4): an array of boxes kept
+        labels (k): an array of labels for each boxes kept
+        probs (k): an array of probabilities for each boxes being in corresponding labels
+    """
+    boxes = boxes[0]
+    confidences = confidences[0]
+    #print(boxes)
+    #print(confidences)
+
+    picked_box_probs = []
+    picked_labels = []
+    for class_index in range(1, confidences.shape[1]):
+        #print(confidences.shape[1])
+        probs = confidences[:, class_index]
+        #print(probs)
+        mask = probs > prob_threshold
+        probs = probs[mask]
+
+        if probs.shape[0] == 0:
+            continue
+        subset_boxes = boxes[mask, :]
+        #print(subset_boxes)
+        box_probs = np.concatenate([subset_boxes, probs.reshape(-1, 1)], axis=1)
+        box_probs = hard_nms(box_probs,
+           iou_threshold=iou_threshold,
+           top_k=top_k,
+           )
+        picked_box_probs.append(box_probs)
+        picked_labels.extend([class_index] * box_probs.shape[0])
+    if not picked_box_probs:
+        return np.array([]), np.array([]), np.array([])
+    picked_box_probs = np.concatenate(picked_box_probs)
+    picked_box_probs[:, 0] *= width
+    picked_box_probs[:, 1] *= height
+    picked_box_probs[:, 2] *= width
+    picked_box_probs[:, 3] *= height
+    return picked_box_probs[:, :4].astype(np.int32), np.array(picked_labels), picked_box_probs[:, 4]