class dl_detection_model():
def __init__(self, model, classNum, confile, confThreshold=0.3, nmsThreshold=0.3, objThreshold=0.5):
try:
# 模型初始设置
self.net = cv2.dnn.readNetFromONNX(model)
self.num_classes = classNum
self.confile = confile
self.confThreshold = confThreshold
self.nmsThreshold = nmsThreshold
self.objThreshold = objThreshold
# 计算设置
self.anchors = [[10, 13, 16, 30, 33, 23], [30, 61, 62, 45, 59, 119], [116, 90, 156, 198, 373, 326]]
self.nl = len(self.anchors)
self.no = self.num_classes + 5
self.grid = [np.zeros(1)] * self.nl
self.stride = np.array([8., 16., 32.])
self.na = len(self.anchors[0]) // 2
self.anchor_grid = np.asarray(self.anchors, dtype=np.float32).reshape(self.nl, 1, -1, 1, 1, 2)
#输出元素
self.outimg = 0
self.outinfo = []
except Exception as e:
LOG.error('Nut recognize model building with exception: {}.'.format(str(e)))
def create_detection(self, image):
try:
self.image = image
# 绘图设置
self.res_rect_size = int((self.image.shape[1] + self.image.shape[0]) / 3000)
self.res_font_size = int((self.image.shape[1] + self.image.shape[0]) / 2400)
self.res_font_thickness = int((self.image.shape[1] + self.image.shape[0]) / 1200)
# 计算
self.dets = self.detectImage()
self.outimg, self.outinfo = self.postprocess(self.dets)
except Exception as e:
LOG.error('Nut recognize model detect with exception: {}.'.format(str(e)))
def get_result_image(self):
return self.outimg
def get_position(self):
return self.outinfo[2], self.outinfo[3]
def get_size(self):
return self.outinfo[4], self.outinfo[5]
def get_classid(self):
return self.outinfo[0]
def get_confidence(self):
return self.outinfo[1]
# 主识别函数
def detectImage(self):
# 推理
blob = cv2.dnn.blobFromImage(self.image, 1 / 255.0, (640, 640), [0, 0, 0], swapRB=True, crop=False)
self.net.setInput(blob)
outs = self.net.forward(self.net.getUnconnectedOutLayersNames())
# 结果变换
z = []
for i in range(self.nl):
bs, _, ny, nx = outs[i].shape
# python的tuple赋值内容后无法修改,需要转换成list才能reshape
outs_listtype = np.array(outs[i])
outs_listtype = outs_listtype.reshape(bs, self.na, self.no, ny, nx).transpose(0, 1, 3, 4, 2)
if self.grid[i].shape[2:4] != outs_listtype.shape[2:4]:
self.grid[i] = self._make_grid(nx, ny)
# sigmoid
y = 1 / (1 + np.exp(-outs_listtype))
# 其实只需要对x,y,w,h做sigmoid变换的, 不过全做sigmoid变换对结果影响不大,因为sigmoid是单调递增函数,那么就不影响类别置信度的排序关系,因此不影响后面的NMS
# 不过设断点查看类别置信度,都是负数,看来有必要做sigmoid变换把概率值强行拉回到0到1的区间内
y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * int(self.stride[i])
y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh
z.append(y.reshape(bs, -1, self.no))
z = np.concatenate(z, axis=1)
return z
##################################################
# 模型结果处理函数
def postprocess(self, outs):
# 绘图准备
frame = self.image
frameHeight = frame.shape[0]
frameWidth = frame.shape[1]
ratioh, ratiow = frameHeight / 640, frameWidth / 640
classIds = []
confidences = []
boxes = []
diresult = []
# 特征提取
for out in outs:
for detection in out:
scores = detection[5:]
classId = np.argmax(scores)
confidence = scores[classId]
if confidence > self.confThreshold and detection[4] > self.objThreshold:
center_x = int(detection[0] * ratiow)
center_y = int(detection[1] * ratioh)
width = int(detection[2] * ratiow)
height = int(detection[3] * ratioh)
left = int(center_x - width / 2)
top = int(center_y - height / 2)
classIds.append(classId)
confidences.append(float(confidence))
boxes.append([left, top, width, height])
# NMS
indices = cv2.dnn.NMSBoxes(boxes, confidences, self.confThreshold, self.nmsThreshold)
# 输出
for i in indices:
box = boxes[i]
left = box[0]
top = box[1]
width = box[2]
height = box[3]
center_x = int(left + width / 2)
center_y = int(top + height / 2)
frame = self.drawPred(frame, classIds[i], confidences[i], left, top, left + width, top + height)
diresult.append((classIds[i], confidences[i], center_x, center_y, width, height))
return frame, diresult
##################################################
# 图像识别框绘制
def drawPred(self, frame, classId, conf, left, top, right, bottom):
# bounding box 绘制
cv2.rectangle(frame, (left, top), (right, bottom), (0, 0, 255), thickness= self.res_rect_size)
# label 绘制
label = '%.2f' % conf
label = '%s:%s' % (self.confile["CheckConfig"]["_class_name_en"][classId], label)
labelSize, baseLine = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
top = max(top, labelSize[1])
# cv.rectangle(frame, (left, top - round(1.5 * labelSize[1])), (left + round(1.5 * labelSize[0]), top + baseLine), (255,255,255), cv.FILLED)
cv2.putText(frame, label, (left, top - 10), cv2.FONT_HERSHEY_SIMPLEX, self.res_font_size, (0, 255, 0), thickness=self.res_font_thickness)
return frame
##################################################
# grid绘制
def _make_grid(self, nx=20, ny=20):
xv, yv = np.meshgrid(np.arange(ny), np.arange(nx))
return np.stack((xv, yv), 2).reshape((1, 1, ny, nx, 2)).astype(np.float32)
opencv dnn +onnx模型 做推理
最新推荐文章于 2025-03-12 10:55:42 发布