| import numpy as np |
| import cv2 |
|
|
| class NanoDet: |
| def __init__(self, modelPath, prob_threshold=0.35, iou_threshold=0.6, backend_id=0, target_id=0): |
| self.strides = (8, 16, 32, 64) |
| self.image_shape = (416, 416) |
| self.reg_max = 7 |
| self.prob_threshold = prob_threshold |
| self.iou_threshold = iou_threshold |
| self.backend_id = backend_id |
| self.target_id = target_id |
| self.project = np.arange(self.reg_max + 1) |
| self.mean = np.array([103.53, 116.28, 123.675], dtype=np.float32).reshape(1, 1, 3) |
| self.std = np.array([57.375, 57.12, 58.395], dtype=np.float32).reshape(1, 1, 3) |
| self.net = cv2.dnn.readNet(modelPath) |
| self.net.setPreferableBackend(self.backend_id) |
| self.net.setPreferableTarget(self.target_id) |
|
|
| self.anchors_mlvl = [] |
| for i in range(len(self.strides)): |
| featmap_size = (int(self.image_shape[0] / self.strides[i]), int(self.image_shape[1] / self.strides[i])) |
| stride = self.strides[i] |
| feat_h, feat_w = featmap_size |
| shift_x = np.arange(0, feat_w) * stride |
| shift_y = np.arange(0, feat_h) * stride |
| xv, yv = np.meshgrid(shift_x, shift_y) |
| xv = xv.flatten() |
| yv = yv.flatten() |
| cx = xv + 0.5 * (stride-1) |
| cy = yv + 0.5 * (stride - 1) |
| |
| anchors = np.column_stack((cx, cy)) |
| self.anchors_mlvl.append(anchors) |
|
|
| @property |
| def name(self): |
| return self.__class__.__name__ |
|
|
| def setBackendAndTarget(self, backendId, targetId): |
| self.backend_id = backendId |
| self.target_id = targetId |
| self.net.setPreferableBackend(self.backend_id) |
| self.net.setPreferableTarget(self.target_id) |
|
|
| def pre_process(self, img): |
| img = img.astype(np.float32) |
| img = (img - self.mean) / self.std |
| blob = cv2.dnn.blobFromImage(img) |
| return blob |
|
|
| def infer(self, srcimg): |
| blob = self.pre_process(srcimg) |
| self.net.setInput(blob) |
| outs = self.net.forward(self.net.getUnconnectedOutLayersNames()) |
| preds = self.post_process(outs) |
| return preds |
|
|
| def post_process(self, preds): |
| cls_scores, bbox_preds = preds[::2], preds[1::2] |
| rescale = False |
| scale_factor = 1 |
| bboxes_mlvl = [] |
| scores_mlvl = [] |
| for stride, cls_score, bbox_pred, anchors in zip(self.strides, cls_scores, bbox_preds, self.anchors_mlvl): |
| if cls_score.ndim==3: |
| cls_score = cls_score.squeeze(axis=0) |
| if bbox_pred.ndim==3: |
| bbox_pred = bbox_pred.squeeze(axis=0) |
|
|
| x_exp = np.exp(bbox_pred.reshape(-1, self.reg_max + 1)) |
| x_sum = np.sum(x_exp, axis=1, keepdims=True) |
| bbox_pred = x_exp / x_sum |
| bbox_pred = np.dot(bbox_pred, self.project).reshape(-1,4) |
| bbox_pred *= stride |
|
|
| nms_pre = 1000 |
| if nms_pre > 0 and cls_score.shape[0] > nms_pre: |
| max_scores = cls_score.max(axis=1) |
| topk_inds = max_scores.argsort()[::-1][0:nms_pre] |
| anchors = anchors[topk_inds, :] |
| bbox_pred = bbox_pred[topk_inds, :] |
| cls_score = cls_score[topk_inds, :] |
|
|
| points = anchors |
| distance = bbox_pred |
| max_shape=self.image_shape |
| x1 = points[:, 0] - distance[:, 0] |
| y1 = points[:, 1] - distance[:, 1] |
| x2 = points[:, 0] + distance[:, 2] |
| y2 = points[:, 1] + distance[:, 3] |
|
|
| if max_shape is not None: |
| x1 = np.clip(x1, 0, max_shape[1]) |
| y1 = np.clip(y1, 0, max_shape[0]) |
| x2 = np.clip(x2, 0, max_shape[1]) |
| y2 = np.clip(y2, 0, max_shape[0]) |
|
|
| |
| bboxes = np.column_stack([x1, y1, x2, y2]) |
| bboxes_mlvl.append(bboxes) |
| scores_mlvl.append(cls_score) |
|
|
| bboxes_mlvl = np.concatenate(bboxes_mlvl, axis=0) |
| if rescale: |
| bboxes_mlvl /= scale_factor |
| scores_mlvl = np.concatenate(scores_mlvl, axis=0) |
| bboxes_wh = bboxes_mlvl.copy() |
| bboxes_wh[:, 2:4] = bboxes_wh[:, 2:4] - bboxes_wh[:, 0:2] |
| classIds = np.argmax(scores_mlvl, axis=1) |
| confidences = np.max(scores_mlvl, axis=1) |
|
|
| indices = cv2.dnn.NMSBoxes(bboxes_wh.tolist(), confidences.tolist(), self.prob_threshold, self.iou_threshold) |
|
|
| if len(indices)>0: |
| det_bboxes = bboxes_mlvl[indices] |
| det_conf = confidences[indices] |
| det_classid = classIds[indices] |
|
|
| return np.concatenate([det_bboxes, det_conf.reshape(-1, 1), det_classid.reshape(-1, 1)], axis=1) |
| else: |
| return np.array([]) |
|
|
