yolov5模型预测

1、模型预测

import onnxruntime
import numpy as np
import cv2
import time


class ChtDeploy():
    def __init__(self, img_path, onnx_path, iou_threshold=0.45, conf_threshold=0.3, detect_w=640, detect_h=640):
        self.img = cv2.imread(img_path)  # h,w,c
        self.img_h = self.img.shape[0]
        self.img_w = self.img.shape[1]
        self.iou_threshold = iou_threshold
        self.conf_threshold = conf_threshold
        self.detect_w = detect_w
        self.detect_h = detect_h
        self.onnx = onnx_path
        self.max_wh = max(self.detect_h, self.detect_w)

    def letterbox(self):
        if (self.img_h == self.detect_h and self.img_w == self.detect_w):
            return self.img

        scale = min(self.detect_w / self.img_w, self.detect_h / self.img_h)  # 缩放比例
        # nw, nh = int(self.img_w * scale), int(self.img_h * scale)
        # image = cv2.resize(self.img, (nw, nh), interpolation=cv2.INTER_LINEAR)
        # img_back = np.ones((self.detect_h, self.detect_w, 3), dtype=np.uint8) * 128
        # # 将image放在画布中心区域-letterbox
        # img_back[(self.detect_h - nh) // 2: (self.detect_h - nh) // 2 + nh, (self.detect_w - nw) // 2:(self.detect_w - nw) // 2 + nw, :] = image

        # 先缩fang，再平移
        h_t, w_t = abs(self.detect_h - scale * self.img_h) / 2, abs(self.detect_w - scale * self.img_w) / 2
        A = np.array([[scale, 0, w_t], [0, scale, h_t]], dtype=np.float32)
        img_back = cv2.warpAffine(self.img, A, (self.detect_w, self.detect_h), borderValue=(128, 128, 128))
        return img_back, A

    def img2input(self, img):
        img = np.transpose(img