|
| 1 | +# -*- coding: utf-8 -*- |
| 2 | +# Copyright (c) Alibaba, Inc. and its affiliates. |
| 3 | +import cv2 |
| 4 | +import numpy as np |
| 5 | + |
| 6 | +import onnxruntime |
| 7 | + |
| 8 | +def nms(boxes, scores, nms_thr): |
| 9 | + """Single class NMS implemented in Numpy.""" |
| 10 | + x1 = boxes[:, 0] |
| 11 | + y1 = boxes[:, 1] |
| 12 | + x2 = boxes[:, 2] |
| 13 | + y2 = boxes[:, 3] |
| 14 | + |
| 15 | + areas = (x2 - x1 + 1) * (y2 - y1 + 1) |
| 16 | + order = scores.argsort()[::-1] |
| 17 | + |
| 18 | + keep = [] |
| 19 | + while order.size > 0: |
| 20 | + i = order[0] |
| 21 | + keep.append(i) |
| 22 | + xx1 = np.maximum(x1[i], x1[order[1:]]) |
| 23 | + yy1 = np.maximum(y1[i], y1[order[1:]]) |
| 24 | + xx2 = np.minimum(x2[i], x2[order[1:]]) |
| 25 | + yy2 = np.minimum(y2[i], y2[order[1:]]) |
| 26 | + |
| 27 | + w = np.maximum(0.0, xx2 - xx1 + 1) |
| 28 | + h = np.maximum(0.0, yy2 - yy1 + 1) |
| 29 | + inter = w * h |
| 30 | + ovr = inter / (areas[i] + areas[order[1:]] - inter) |
| 31 | + |
| 32 | + inds = np.where(ovr <= nms_thr)[0] |
| 33 | + order = order[inds + 1] |
| 34 | + |
| 35 | + return keep |
| 36 | + |
| 37 | +def multiclass_nms(boxes, scores, nms_thr, score_thr): |
| 38 | + """Multiclass NMS implemented in Numpy. Class-aware version.""" |
| 39 | + final_dets = [] |
| 40 | + num_classes = scores.shape[1] |
| 41 | + for cls_ind in range(num_classes): |
| 42 | + cls_scores = scores[:, cls_ind] |
| 43 | + valid_score_mask = cls_scores > score_thr |
| 44 | + if valid_score_mask.sum() == 0: |
| 45 | + continue |
| 46 | + else: |
| 47 | + valid_scores = cls_scores[valid_score_mask] |
| 48 | + valid_boxes = boxes[valid_score_mask] |
| 49 | + keep = nms(valid_boxes, valid_scores, nms_thr) |
| 50 | + if len(keep) > 0: |
| 51 | + cls_inds = np.ones((len(keep), 1)) * cls_ind |
| 52 | + dets = np.concatenate( |
| 53 | + [valid_boxes[keep], valid_scores[keep, None], cls_inds], 1 |
| 54 | + ) |
| 55 | + final_dets.append(dets) |
| 56 | + if len(final_dets) == 0: |
| 57 | + return None |
| 58 | + return np.concatenate(final_dets, 0) |
| 59 | + |
| 60 | +def demo_postprocess(outputs, img_size, p6=False): |
| 61 | + grids = [] |
| 62 | + expanded_strides = [] |
| 63 | + strides = [8, 16, 32] if not p6 else [8, 16, 32, 64] |
| 64 | + |
| 65 | + hsizes = [img_size[0] // stride for stride in strides] |
| 66 | + wsizes = [img_size[1] // stride for stride in strides] |
| 67 | + |
| 68 | + for hsize, wsize, stride in zip(hsizes, wsizes, strides): |
| 69 | + xv, yv = np.meshgrid(np.arange(wsize), np.arange(hsize)) |
| 70 | + grid = np.stack((xv, yv), 2).reshape(1, -1, 2) |
| 71 | + grids.append(grid) |
| 72 | + shape = grid.shape[:2] |
| 73 | + expanded_strides.append(np.full((*shape, 1), stride)) |
| 74 | + |
| 75 | + grids = np.concatenate(grids, 1) |
| 76 | + expanded_strides = np.concatenate(expanded_strides, 1) |
| 77 | + outputs[..., :2] = (outputs[..., :2] + grids) * expanded_strides |
| 78 | + outputs[..., 2:4] = np.exp(outputs[..., 2:4]) * expanded_strides |
| 79 | + |
| 80 | + return outputs |
| 81 | + |
| 82 | +def preprocess(img, input_size, swap=(2, 0, 1)): |
| 83 | + if len(img.shape) == 3: |
| 84 | + padded_img = np.ones((input_size[0], input_size[1], 3), dtype=np.uint8) * 114 |
| 85 | + else: |
| 86 | + padded_img = np.ones(input_size, dtype=np.uint8) * 114 |
| 87 | + |
| 88 | + r = min(input_size[0] / img.shape[0], input_size[1] / img.shape[1]) |
| 89 | + resized_img = cv2.resize( |
| 90 | + img, |
| 91 | + (int(img.shape[1] * r), int(img.shape[0] * r)), |
| 92 | + interpolation=cv2.INTER_LINEAR, |
| 93 | + ).astype(np.uint8) |
| 94 | + padded_img[: int(img.shape[0] * r), : int(img.shape[1] * r)] = resized_img |
| 95 | + |
| 96 | + padded_img = padded_img.transpose(swap) |
| 97 | + padded_img = np.ascontiguousarray(padded_img, dtype=np.float32) |
| 98 | + return padded_img, r |
| 99 | + |
| 100 | +def inference_detector(session, oriImg): |
| 101 | + input_shape = (640,640) |
| 102 | + img, ratio = preprocess(oriImg, input_shape) |
| 103 | + |
| 104 | + ort_inputs = {session.get_inputs()[0].name: img[None, :, :, :]} |
| 105 | + output = session.run(None, ort_inputs) |
| 106 | + predictions = demo_postprocess(output[0], input_shape)[0] |
| 107 | + |
| 108 | + boxes = predictions[:, :4] |
| 109 | + scores = predictions[:, 4:5] * predictions[:, 5:] |
| 110 | + |
| 111 | + boxes_xyxy = np.ones_like(boxes) |
| 112 | + boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2]/2. |
| 113 | + boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3]/2. |
| 114 | + boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2]/2. |
| 115 | + boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3]/2. |
| 116 | + boxes_xyxy /= ratio |
| 117 | + dets = multiclass_nms(boxes_xyxy, scores, nms_thr=0.45, score_thr=0.1) |
| 118 | + if dets is not None: |
| 119 | + final_boxes, final_scores, final_cls_inds = dets[:, :4], dets[:, 4], dets[:, 5] |
| 120 | + isscore = final_scores>0.3 |
| 121 | + iscat = final_cls_inds == 0 |
| 122 | + isbbox = [ i and j for (i, j) in zip(isscore, iscat)] |
| 123 | + final_boxes = final_boxes[isbbox] |
| 124 | + else: |
| 125 | + final_boxes = np.array([]) |
| 126 | + |
| 127 | + return final_boxes |
0 commit comments