Python实现ONNXRuntime推理YOLOv11模型

Python实现ONNXRuntime推理YOLOv11模型,主要在于onnxruntime推理后的后处理部分

1、安装依赖

pip install opencv-python onnxruntime numpy

2、ONNX模型导出(可选)

from ultralytics import YOLO
 
# Load a model
model = YOLO("yolo11n.pt")  # load an official model
model = YOLO("best.pt")  # load a custom trained model
 
# Export the model
model.export(format="onnx", simplify=True, half=True)

3、Python推理代码

import cv2
import numpy as np
import onnxruntime as ort
from math import exp
 
# 常量配置
CLASSES = ['class1']  # 模型类别

np.random.seed(1)
COLORS = np.random.uniform(0, 255, size=(len(CLASSES), 3))  # 随机颜色
meshgrid = []
class_num = len(CLASSES)
headNum = 3
strides = [8, 16, 32]
mapSize = [[80, 80], [40, 40], [20, 20]]
input_imgH = 640
input_imgW = 640
 
 
class DetectBox:
    """检测框类"""
    def __init__(self, classId, score, xmin, ymin, xmax, ymax):
        self.classId = classId
        self.score = score
        self.xmin = xmin
        self.ymin = ymin
        self.xmax = xmax
        self.ymax = ymax
 

 
class YOLODetector:
    def __init__(self, model_path='./yolov11n.onnx', conf_thresh=0.5, iou_thresh=0.45):
        self.model_path = model_path
        self.conf_thresh = conf_thresh
        self.iou_thresh = iou_thresh
        self.ort_session = ort.InferenceSession(self.model_path)
        self.generate_meshgrid()
 
    @staticmethod
    def sigmoid(x):
        return 1 / (1 + exp(-x))
 
    @staticmethod
    def preprocess_image(img_src, resize_w, resize_h):
        image = cv2.resize(img_src, (resize_w, resize_h), interpolation=cv2.INTER_LINEAR)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        image = image.astype(np.float32)
        image /= 255.0
        return image
 
    def generate_meshgrid(self):
        for index in range(headNum):
            for i in range(mapSize[index][0]):
                for j in range(mapSize[index][1]):
                    meshgrid.append(j + 0.5)
                    meshgrid.append(i + 0.5)
 
    def iou(self, xmin1, ymin1, xmax1, ymax1, xmin2, ymin2, xmax2, ymax2):
        xmin = max(xmin1, xmin2)
        ymin = max(ymin1, ymin2)
        xmax = min(xmax1, xmax2)
        ymax = min(ymax1, ymax2)
 
        innerWidth = max(0, xmax - xmin)
        innerHeight = max(0, ymax - ymin)
 
        innerArea = innerWidth * innerHeight
        area1 = (xmax1 - xmin1) * (ymax1 - ymin1)
        area2 = (xmax2 - xmin2) * (ymax2 - ymin2)
        total = area1 + area2 - innerArea
 
        return innerArea / total
 
    def nms(self, detectResult):
        predBoxs = []
        sort_detectboxs = sorted(detectResult, key=lambda x: x.score, reverse=True)
 
        for i in range(len(sort_detectboxs)):
            if sort_detectboxs[i].classId != -1:
                predBoxs.append(sort_detectboxs[i])
                for j in range(i + 1, len(sort_detectboxs), 1):
                    if sort_detectboxs[i].classId == sort_detectboxs[j].classId:
                        iou = self.iou(
                            sort_detectboxs[i].xmin, sort_detectboxs[i].ymin,
                            sort_detectboxs[i].xmax, sort_detectboxs[i].ymax,
                            sort_detectboxs[j].xmin, sort_detectboxs[j].ymin,
                            sort_detectboxs[j].xmax, sort_detectboxs[j].ymax
                        )
                        if iou > self.iou_thresh:
                            sort_detectboxs[j].classId = -1
        return predBoxs
 
    def postprocess(self, out, img_h, img_w):
        detectResult = []
        output = out[0][0]  # 去掉 batch 维度,变为 (5 + num_classes, num_boxes)
 
        # 提取预测框信息
        reg = output[0:4, :]  # 回归框的 x, y, w, h
        conf = output[4, :]  # 置信度
 
        # 检查是否是多类别模型
        if output.shape[0] > 5:  # 如果输出维度大于 5,说明是多类别模型
            class_probs = output[5:, :]  # 类别概率 (num_classes, num_boxes)
            is_multiclass = True
        else:
            is_multiclass = False
 
        scale_h = img_h / input_imgH
        scale_w = img_w / input_imgW
 
        for i in range(reg.shape[1]):  # 遍历所有预测框
            x, y, w, h = reg[:, i]
            score = self.sigmoid(conf[i])  # 使用 sigmoid 激活置信度
            # score = float(conf[i])  # 使用 sigmoid 激活置信度
 
            if is_multiclass:  # 多类别模型
                # 使用 softmax 激活类别概率
                class_prob = np.exp(class_probs[:, i]) / np.sum(np.exp(class_probs[:, i]))
                class_id = np.argmax(class_prob)  # 获取最大概率的类别索引
                class_score = class_prob[class_id]  # 获取该类别的概率
 
                # 综合置信度和类别概率
                final_score = score * class_score
            else:  # 单类别模型
                class_id = 0  # 单类别情况下,类别索引固定为 0
                final_score = score  # 置信度即为最终得分
 
            if final_score > self.conf_thresh:  # 过滤低置信度框
                xmin = max(0, (x - w / 2) * scale_w)
                ymin = max(0, (y - h / 2) * scale_h)
                xmax = min(img_w, (x + w / 2) * scale_w)
                ymax = min(img_h, (y + h / 2) * scale_h)
 
                # 添加框信息
                box = DetectBox(classId=class_id, score=final_score, xmin=xmin, ymin=ymin, xmax=xmax, ymax=ymax)
                detectResult.append(box)
 
        predBox = self.nms(detectResult)  # 非极大值抑制
        return predBox
 
    def detect(self, img_path):
        if isinstance(img_path, str):
            orig = cv2.imread(img_path)
        else:
            orig = img_path
 
        img_h, img_w = orig.shape[:2]
        image = self.preprocess_image(orig, input_imgW, input_imgH)
        image = image.transpose((2, 0, 1))
        image = np.expand_dims(image, axis=0)
 
        pred_results = self.ort_session.run(None, {'images': image})
 
        # 打印模型输出形状(调试用)
        print(f"Model output shape: {pred_results[0].shape}")
 
        predbox = self.postprocess(pred_results, img_h, img_w)
 
        boxes = []
        scores = []
        class_ids = []
 
        for box in predbox:
            boxes.append([int(box.xmin), int(box.ymin), int(box.xmax), int(box.ymax)])
            scores.append(box.score)
            class_ids.append(box.classId)
 
        return boxes, scores, class_ids
 
    def draw_detections(self, image, boxes, scores, class_ids, mask_alpha=0.3):
        """
        Combines drawing masks, boxes, and text annotations on detected objects.
        Parameters:
        - image: Input image.
        - boxes: Array of bounding boxes.
        - scores: Confidence scores for each detected object.
        - class_ids: Detected object class IDs.
        - mask_alpha: Transparency of the mask overlay.
        """
        det_img = image.copy()
 
        img_height, img_width = image.shape[:2]
        font_size = min([img_height, img_width]) * 0.001
        text_thickness = int(min([img_height, img_width]) * 0.001)
 
        mask_img = image.copy()
 
        # Draw bounding boxes, masks, and text annotations
        for class_id, box, score in zip(class_ids, boxes, scores):
            color = COLORS[class_id]
            x1, y1, x2, y2 = box[0], box[1], box[2], box[3]
 
            # Draw fill rectangle for mask
            cv2.rectangle(mask_img, (x1, y1), (x2, y2), color, -1)
 
            # Draw bounding box
            cv2.rectangle(det_img, (x1, y1), (x2, y2), color, 2)
 
            # Prepare text (label and score)
            label = CLASSES[class_id]
            caption = f'{label} {int(score * 100)}%'
 
            # Calculate text size and position
            (tw, th), _ = cv2.getTextSize(text=caption, fontFace=cv2.FONT_HERSHEY_SIMPLEX,
                                          fontScale=font_size, thickness=text_thickness)
            th = int(th * 1.2)
 
            # Draw filled rectangle for text background
            cv2.rectangle(det_img, (x1, y1), (x1 + tw, y1 - th), color, -1)
 
            # Draw text over the filled rectangle
            cv2.putText(det_img, caption, (x1, y1), cv2.FONT_HERSHEY_SIMPLEX, font_size,
                        (255, 255, 255), text_thickness, cv2.LINE_AA)
 
        # Blend the mask image with the original image
        det_img = cv2.addWeighted(mask_img, mask_alpha, det_img, 1 - mask_alpha, 0)
 
        return det_img
    
if __name__ == "__main__":
    
    model_path = './models/epoch_100_N.onnx'
    img_path = "/mnt/d/data/001/SD-001-00920/SI-001-01708.png"
 
    detector = YOLODetector(model_path=model_path, conf_thresh=0.55, iou_thresh=0.7)
 
    boxes, scores, class_ids = detector.detect(img_path)
    print("检测到的框:", boxes)
    print("检测到的分数:", scores)
    print("检测到的类别:", class_ids)
 
    image = cv2.imread(img_path)
    result_img = detector.draw_detections(image, boxes, scores, class_ids)
    cv2.imwrite('result.jpg', result_img)
    cv2.imshow('Detection Results', result_img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

4、yolo11架构图

在这里插入图片描述

### 使用 ONNX Runtime 进行 YOLOv5 模型推理 为了使用 ONNX Runtime 对 YOLOv5 模型进行高效推理,可以遵循以下方法: #### 准备工作 确保安装了必要的依赖项。可以通过 pip 安装 `onnx` 和 `onnxruntime` 库。 ```bash pip install onnx onnxruntime ``` #### 导出 PyTorch 模型ONNX 格式 将训练好的 `.pt` 文件转换成 `.onnx` 文件以便于在 ONNX Runtime 中加载和运行。这一步骤对于提高性能至关重要[^1]。 ```python import torch.onnx from models.experimental import attempt_load model = attempt_load('path/to/yolov5s.pt', map_location=torch.device('cpu')) # 加载模型 dummy_input = torch.randn(1, 3, 640, 640) # 创建虚拟输入张量 input_names = ["image"] output_names = ['output'] torch.onnx.export( model, dummy_input, "yolov5s.onnx", input_names=input_names, output_names=output_names, opset_version=12 ) ``` #### 使用 ONNX Runtime 执行推理 一旦拥有了 `.onnx` 文件,则可以在 Python 或其他支持的语言环境中利用 ONNX Runtime 来执行高效的推理操作。 ```python import numpy as np import cv2 import onnxruntime as ort session = ort.InferenceSession("yolov5s.onnx") def preprocess(image_path): image = cv2.imread(image_path) image_resized = cv2.resize(image, (640, 640)) blob = cv2.dnn.blobFromImage(image_resized, scalefactor=1 / 255.0, size=(640, 640), swapRB=True, crop=False) return blob blob = preprocess("bus.jpg") outputs = session.run(None, {"image": blob}) print(outputs) ``` 上述代码展示了如何准备图像数据并通过调用 `InferenceSession.run()` 方法来获取预测结果。需要注意的是,在实际应用中可能还需要处理输出以提取有用的信息,比如边界框的位置、类别标签以及置信度分数等。
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

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

抵扣说明:

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

余额充值