opencv dnn +onnx模型 做推理

该代码定义了一个名为`classdl_detection_model`的类,用于处理ONNX模型进行目标检测。初始化方法设置模型参数,包括置信度阈值等。`detectImage`方法执行模型推理,`postprocess`方法处理结果,包括NMS非极大抑制。此外,还有绘制检测框的辅助方法。

摘要生成于 C知道 ,由 DeepSeek-R1 满血版支持, 前往体验 >

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变换把概率值强行拉回到01的区间内
            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)
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值