Spaces:
Runtime error
Runtime error
| #!/usr/bin/env python3 | |
| # -*- coding:utf-8 -*- | |
| # Copyright (c) 2014-2021 Megvii Inc. All rights reserved. | |
| import numpy as np | |
| import os | |
| __all__ = ["mkdir", "nms", "multiclass_nms", "demo_postprocess"] | |
| def mkdir(path): | |
| if not os.path.exists(path): | |
| os.makedirs(path) | |
| def nms(boxes, scores, nms_thr): | |
| """Single class NMS implemented in Numpy.""" | |
| x1 = boxes[:, 0] | |
| y1 = boxes[:, 1] | |
| x2 = boxes[:, 2] | |
| y2 = boxes[:, 3] | |
| areas = (x2 - x1 + 1) * (y2 - y1 + 1) | |
| order = scores.argsort()[::-1] | |
| keep = [] | |
| while order.size > 0: | |
| i = order[0] | |
| keep.append(i) | |
| 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:]]) | |
| w = np.maximum(0.0, xx2 - xx1 + 1) | |
| h = np.maximum(0.0, yy2 - yy1 + 1) | |
| inter = w * h | |
| ovr = inter / (areas[i] + areas[order[1:]] - inter) | |
| inds = np.where(ovr <= nms_thr)[0] | |
| order = order[inds + 1] | |
| return keep | |
| def multiclass_nms(boxes, scores, nms_thr, score_thr): | |
| """Multiclass NMS implemented in Numpy""" | |
| final_dets = [] | |
| num_classes = scores.shape[1] | |
| for cls_ind in range(num_classes): | |
| cls_scores = scores[:, cls_ind] | |
| valid_score_mask = cls_scores > score_thr | |
| if valid_score_mask.sum() == 0: | |
| continue | |
| else: | |
| valid_scores = cls_scores[valid_score_mask] | |
| valid_boxes = boxes[valid_score_mask] | |
| keep = nms(valid_boxes, valid_scores, nms_thr) | |
| if len(keep) > 0: | |
| cls_inds = np.ones((len(keep), 1)) * cls_ind | |
| dets = np.concatenate( | |
| [valid_boxes[keep], valid_scores[keep, None], cls_inds], 1 | |
| ) | |
| final_dets.append(dets) | |
| if len(final_dets) == 0: | |
| return None | |
| return np.concatenate(final_dets, 0) | |
| def demo_postprocess(outputs, img_size, p6=False): | |
| grids = [] | |
| expanded_strides = [] | |
| if not p6: | |
| strides = [8, 16, 32] | |
| else: | |
| strides = [8, 16, 32, 64] | |
| hsizes = [img_size[0] // stride for stride in strides] | |
| wsizes = [img_size[1] // stride for stride in strides] | |
| for hsize, wsize, stride in zip(hsizes, wsizes, strides): | |
| xv, yv = np.meshgrid(np.arange(wsize), np.arange(hsize)) | |
| grid = np.stack((xv, yv), 2).reshape(1, -1, 2) | |
| grids.append(grid) | |
| shape = grid.shape[:2] | |
| expanded_strides.append(np.full((*shape, 1), stride)) | |
| grids = np.concatenate(grids, 1) | |
| expanded_strides = np.concatenate(expanded_strides, 1) | |
| outputs[..., :2] = (outputs[..., :2] + grids) * expanded_strides | |
| outputs[..., 2:4] = np.exp(outputs[..., 2:4]) * expanded_strides | |
| return outputs | |