|
|
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([]) |
|
|
|
