File size: 4,874 Bytes
da19665 a07f7bd d2b2b68 da19665 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 |
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.stack((cx, cy), axis=-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.stack([x1, y1, x2, y2], axis=-1)
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([])
|