onnx/ultraface

1# SPDX-License-Identifier: MIT2 3import numpy as np4 5def area_of(left_top, right_bottom):6    """7    Compute the areas of rectangles given two corners.8    Args:9        left_top (N, 2): left top corner.10        right_bottom (N, 2): right bottom corner.11    Returns:12        area (N): return the area.13    """14    hw = np.clip(right_bottom - left_top, 0.0, None)15    return hw[..., 0] * hw[..., 1]16 17def iou_of(boxes0, boxes1, eps=1e-5):18    """19    Return intersection-over-union (Jaccard index) of boxes.20    Args:21        boxes0 (N, 4): ground truth boxes.22        boxes1 (N or 1, 4): predicted boxes.23        eps: a small number to avoid 0 as denominator.24    Returns:25        iou (N): IoU values.26    """27    overlap_left_top = np.maximum(boxes0[..., :2], boxes1[..., :2])28    overlap_right_bottom = np.minimum(boxes0[..., 2:], boxes1[..., 2:])29 30    overlap_area = area_of(overlap_left_top, overlap_right_bottom)31    area0 = area_of(boxes0[..., :2], boxes0[..., 2:])32    area1 = area_of(boxes1[..., :2], boxes1[..., 2:])33    return overlap_area / (area0 + area1 - overlap_area + eps)34 35def hard_nms(box_scores, iou_threshold, top_k=-1, candidate_size=200):36    """37    Perform hard non-maximum-supression to filter out boxes with iou greater38    than threshold39    Args:40        box_scores (N, 5): boxes in corner-form and probabilities.41        iou_threshold: intersection over union threshold.42        top_k: keep top_k results. If k <= 0, keep all the results.43        candidate_size: only consider the candidates with the highest scores.44    Returns:45        picked: a list of indexes of the kept boxes46    """47    scores = box_scores[:, -1]48    boxes = box_scores[:, :-1]49    picked = []50    indexes = np.argsort(scores)51    indexes = indexes[-candidate_size:]52    while len(indexes) > 0:53        current = indexes[-1]54        picked.append(current)55        if 0 < top_k == len(picked) or len(indexes) == 1:56            break57        current_box = boxes[current, :]58        indexes = indexes[:-1]59        rest_boxes = boxes[indexes, :]60        iou = iou_of(61            rest_boxes,62            np.expand_dims(current_box, axis=0),63        )64        indexes = indexes[iou <= iou_threshold]65 66    return box_scores[picked, :]67 68def predict(width, height, confidences, boxes, prob_threshold, iou_threshold=0.5, top_k=-1):69    """70    Select boxes that contain human faces71    Args:72        width: original image width73        height: original image height74        confidences (N, 2): confidence array75        boxes (N, 4): boxes array in corner-form76        iou_threshold: intersection over union threshold.77        top_k: keep top_k results. If k <= 0, keep all the results.78    Returns:79        boxes (k, 4): an array of boxes kept80        labels (k): an array of labels for each boxes kept81        probs (k): an array of probabilities for each boxes being in corresponding labels82    """83    boxes = boxes[0]84    confidences = confidences[0]85    #print(boxes)86    #print(confidences)87 88    picked_box_probs = []89    picked_labels = []90    for class_index in range(1, confidences.shape[1]):91        #print(confidences.shape[1])92        probs = confidences[:, class_index]93        #print(probs)94        mask = probs > prob_threshold95        probs = probs[mask]96 97        if probs.shape[0] == 0:98            continue99        subset_boxes = boxes[mask, :]100        #print(subset_boxes)101        box_probs = np.concatenate([subset_boxes, probs.reshape(-1, 1)], axis=1)102        box_probs = hard_nms(box_probs,103           iou_threshold=iou_threshold,104           top_k=top_k,105           )106        picked_box_probs.append(box_probs)107        picked_labels.extend([class_index] * box_probs.shape[0])108    if not picked_box_probs:109        return np.array([]), np.array([]), np.array([])110    picked_box_probs = np.concatenate(picked_box_probs)111    picked_box_probs[:, 0] *= width112    picked_box_probs[:, 1] *= height113    picked_box_probs[:, 2] *= width114    picked_box_probs[:, 3] *= height115    return picked_box_probs[:, :4].astype(np.int32), np.array(picked_labels), picked_box_probs[:, 4]