Yolo26 模型转换 onnx 再转换 om 模型到 Ascend310B4 运行经过

Yolo26 模型转换 onnx 再转换 om 模型到 Ascend310B4 运行经过

前言

原来以为说什么 yolo26 转换成 om 模型不能在昇腾上面运行，实际呢就是把训练好的模型转换好就行了，这里只要昇腾芯片对 onnx 模型版本的要求（一般 opset = 11 较为稳定），以及昇腾芯片能够适配的 CANN 版本就好了，其他就是一个运行匹配的问题，此次为什么能够花费这么久呢？完全是因为自己不理解 yolo26 模型的输出是什么意思，蒙起眼睛来打靶，完全打不中，后来参考了一个博主的，把需求弄清楚，就好了。

版本对应问题

这里再次强调一下版本

我的芯片是 Ascend310B4 对应的 HDK为23.0.rc3

芯片适配的CANN版本是 8.1.RC1 查询：CANN版本兼容性-CANN社区版8.1.RC1-昇腾社区

适配的opset版本是 11 查询：通过查看昇腾 CANN 对 ONNX 的支持 链接====>https://www.hiascend.com/document/detail/zh/canncommercial/900/API/aolapi/operatorlist_00177.html 通过查看 ONNX 对应的 ONNX opset 链接 ====> https://runtime.onnx.org.cn/docs/reference/compatibility.html

代码

对于代码来说，主要是一个理解转化的模型需要进行 NMS 就好了，理解模型的输出，就能够很好的解决问题。

比如模型的输入是[batch,300,6]那么就要明白，这个数据代表的意思是什么？

加入 batch 是 1 ，那么如果输入一张图片的话，就会有 300 个框的结果，6 代表的是 x1, y1, x2, y2, score, class_id

参数	含义
x1	左上角 x 坐标
y1	左上角 y 坐标
x2	右下角 x 坐标
y2	右下角 y 坐标

score = 模型认为这个框里"确实有目标"的概率

class_id = 表示这个目标属于哪个类别

总结成一句话就是：在图像的某个位置，我认为那里有一个目标，它的置信度是多少，它属于哪个类别

结果很明确，一张图中不太可能有 300 个目标，那么我们就需要对结果做筛选，选出符合要求的目标，此时就需要做 NMS 非极大值抑制，非极大值抑制是关键，因为 yolo26 的模型导出成 onnx 并没有直接给你抑制

代码如下：

from ais_bench.infer.interface import InferSession  # 导入Ascend推理接口
import cv2  # 导入OpenCV库，用于图像处理
import numpy as np  # 导入NumPy库，用于数值计算
from pathlib import Path  # 导入Path类，用于路径操作

# =========================
# 类别定义
# =========================
CLASS_NAMES = {
    0: "frogman",  # 类别ID 0: 蛙人
    1: "auv",      # 类别ID 1: 自主水下航行器
    2: "chain",    # 类别ID 2: 链条
    3: "boat",     # 类别ID 3: 船只
    4: "cage",     # 类别ID 4: 笼子
    5: "ball"      # 类别ID 5: 球
}

IMG_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff", ".webp"}  # 支持的图像文件扩展名集合


# =========================
# Letterbox
# =========================
def letterbox(img, new_shape=640, color=(114, 114, 114)):
    """
    将图像通过letterbox方式缩放到指定大小，保持宽高比并填充边缘
    
    参数:
        img: 输入图像
        new_shape: 目标尺寸（宽高相同），默认640
        color: 填充颜色，默认灰色(114,114,114)
    
    返回:
        img: 处理后的图像
        r: 缩放比例
        (dw, dh): 填充的宽度和高度（未取整前的）
    """
    h, w = img.shape[:2]  # 获取原始图像高度和宽度
    r = min(new_shape / h, new_shape / w)  # 计算缩放比例，选择较小的缩放因子以确保图像能完整放入

    new_w, new_h = int(w * r), int(h * r)  # 计算缩放后的新宽度和新高度
    resized = cv2.resize(img, (new_w, new_h))  # 将图像缩放到新尺寸

    dw = new_shape - new_w  # 计算需要填充的宽度差值
    dh = new_shape - new_h  # 计算需要填充的高度差值
    dw /= 2  # 左右均分填充宽度
    dh /= 2  # 上下均分填充高度

    top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))  # 计算上边和下边的填充像素数
    left, right = int(round(dw - 0.1)), int(round(dw + 0.1))  # 计算左边和右边的填充像素数

    img = cv2.copyMakeBorder(  # 添加边框填充
        resized,  # 输入图像
        top, bottom, left, right,  # 四边填充大小
        cv2.BORDER_CONSTANT,  # 边框类型为常数填充
        value=color  # 填充颜色值
    )

    return img, r, (dw, dh)  # 返回处理后的图像、缩放比例和填充值


# =========================
# NMS
# =========================
def nms(boxes, scores, iou_thres=0.45):
    """
    非极大值抑制算法，去除重叠度高的冗余检测框
    
    参数:
        boxes: 边界框坐标列表 [[x1,y1,x2,y2], ...]
        scores: 对应的置信度分数
        iou_thres: IoU阈值，高于此阈值的框会被抑制
    
    返回:
        保留的边界框索引列表
    """
    if len(boxes) == 0:  # 如果没有边界框，直接返回空列表
        return []

    xywh = []  # 存储转换为[x,y,width,height]格式的边界框
    for x1, y1, x2, y2 in boxes:  # 遍历所有边界框
        xywh.append([x1, y1, x2 - x1, y2 - y1])  # 转换为[x,y,宽度,高度]格式

    idxs = cv2.dnn.NMSBoxes(  # 调用OpenCV的NMS函数
        xywh,  # 边界框列表
        scores.tolist(),  # 置信度分数列表
        0.0,  # 置信度阈值（此处设为0，因为之前已经过滤过）
        iou_thres  # IoU阈值
    )

    if len(idxs) == 0:  # 如果没有保留的框
        return []

    return idxs.flatten()  # 将索引展平为一维数组返回


# =========================
# 后处理（YOLOv26: 1,300,6）
# =========================
def postprocess(pred, orig_shape, ratio, pad, conf_thres=0.25, iou_thres=0.45):
    """
    模型输出后处理：置信度过滤、NMS、坐标还原
    
    参数:
        pred: 模型预测输出，形状为(300,6)，每行格式[x1,y1,x2,y2,score,class_id]
        orig_shape: 原始图像尺寸(高度,宽度)
        ratio: 缩放比例
        pad: 填充值(dw,dh)
        conf_thres: 置信度阈值
        iou_thres: NMS的IoU阈值
    
    返回:
        boxes: 过滤并还原后的边界框
        scores: 对应的置信度分数
        cls_ids: 对应的类别ID
    """
    orig_h, orig_w = orig_shape  # 获取原始图像高度和宽度
    dw, dh = pad  # 获取填充值

    pred = np.squeeze(pred)  # 去除批次维度，从(1,300,6)变为(300,6)

    boxes = pred[:, :4]  # 提取边界框坐标 [x1,y1,x2,y2]
    scores = pred[:, 4]  # 提取置信度分数
    cls_ids = pred[:, 5].astype(int)  # 提取类别ID并转换为整数类型

    # 置信度过滤：保留置信度大于阈值的检测结果
    mask = scores > conf_thres
    boxes, scores, cls_ids = boxes[mask], scores[mask], cls_ids[mask]

    if len(boxes) == 0:  # 如果没有检测结果
        return [], [], []

    # 非极大值抑制(NMS)
    keep = nms(boxes, scores, iou_thres)

    boxes = boxes[keep]  # 保留NMS后的边界框
    scores = scores[keep]  # 保留对应的置信度
    cls_ids = cls_ids[keep]  # 保留对应的类别ID

    # 还原坐标：将letterbox后的坐标映射回原始图像坐标
    boxes[:, [0, 2]] = (boxes[:, [0, 2]] - dw) / ratio  # 还原x1和x2坐标
    boxes[:, [1, 3]] = (boxes[:, [1, 3]] - dh) / ratio  # 还原y1和y2坐标

    # 坐标裁剪：确保边界框不超出原始图像范围
    boxes[:, [0, 2]] = boxes[:, [0, 2]].clip(0, orig_w)  # 裁剪x坐标到[0,宽度]
    boxes[:, [1, 3]] = boxes[:, [1, 3]].clip(0, orig_h)  # 裁剪y坐标到[0,高度]

    return boxes, scores, cls_ids  # 返回处理后的结果


# =========================
# 画框（已修复OpenCV报错）
# =========================
def draw_boxes(img, boxes, scores, cls_ids):
    """
    在图像上绘制检测框、标签和置信度
    
    参数:
        img: 输入图像
        boxes: 边界框坐标列表
        scores: 置信度分数列表
        cls_ids: 类别ID列表
    
    返回:
        绘制了检测结果的图像
    """
    for box, score, cid in zip(boxes, scores, cls_ids):  # 遍历每个检测结果
        x1, y1, x2, y2 = map(int, box)  # 将坐标转换为整数类型

        label = f"{CLASS_NAMES.get(int(cid), cid)} {score:.2f}"  # 生成标签文本：类别名+置信度

        # ⚠️ 关键修复：必须强制 int（避免 OpenCV 报错）
        # 根据类别ID动态生成颜色（确保值在0-255范围内且为整数）
        color = (
            int((37 * int(cid)) % 255),  # 红色分量
            int((17 * int(cid)) % 255),  # 绿色分量
            int((29 * int(cid)) % 255)   # 蓝色分量
        )

        # 绘制边界框矩形
        cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)  # 线条粗细为2像素

        # 获取标签文本的尺寸
        (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)

        # 计算标签背景矩形的上边界（避免超出图像顶部）
        y_top = max(0, y1 - th - 6)
        # 绘制标签背景矩形
        cv2.rectangle(img, (x1, y_top), (x1 + tw, y1), color, -1)  # -1表示填充矩形

        # 绘制标签文本
        cv2.putText(
            img,  # 目标图像
            label,  # 文本内容
            (x1, max(0, y1 - 4)),  # 文本位置
            cv2.FONT_HERSHEY_SIMPLEX,  # 字体类型
            0.6,  # 字体大小
            (255, 255, 255),  # 文本颜色（白色）
            2,  # 文本粗细
            cv2.LINE_AA  # 抗锯齿线条类型
        )

    return img  # 返回绘制后的图像


# =========================
# 获取图片
# =========================
def get_images(folder: Path):
    """
    递归获取文件夹下所有支持的图像文件
    
    参数:
        folder: 文件夹路径（Path对象）
    
    返回:
        图像文件路径列表
    """
    imgs = []  # 存储图像路径的列表
    for p in sorted(folder.rglob("*")):  # 递归遍历文件夹下所有文件（排序后）
        if p.suffix.lower() in IMG_EXTS:  # 如果文件扩展名在支持的图像格式集合中
            imgs.append(p)  # 添加到列表
    return imgs


# =========================
# 主函数
# =========================
def main():
    """主函数：解析命令行参数、加载模型、执行推理并保存结果"""
    import argparse  # 导入命令行参数解析模块

    parser = argparse.ArgumentParser()  # 创建参数解析器
    parser.add_argument("--model", type=str, required=True)  # 模型文件路径（必需）
    parser.add_argument("--source", type=str, required=True)  # 输入图像源文件夹路径（必需）
    parser.add_argument("--save", type=str, default="./output")  # 输出保存目录，默认为./output
    parser.add_argument("--device", type=int, default=0)  # 昇腾设备ID，默认为0
    parser.add_argument("--conf", type=float, default=0.20)  # 置信度阈值，默认0.20
    parser.add_argument("--iou", type=float, default=0.40)  # NMS的IoU阈值，默认0.40

    args = parser.parse_args()  # 解析命令行参数

    # =========================
    # load model
    # =========================
    session = InferSession(args.device, args.model)  # 创建推理会话，加载模型到指定设备

    save_dir = Path(args.save)  # 创建输出保存目录的Path对象
    img_dir = save_dir / "images"  # 保存绘制后图像的子目录
    label_dir = save_dir / "labels"  # 保存YOLO格式标签的子目录

    img_dir.mkdir(parents=True, exist_ok=True)  # 创建图像保存目录（如果父目录不存在则递归创建）
    label_dir.mkdir(parents=True, exist_ok=True)  # 创建标签保存目录

    # =========================
    # images
    # =========================
    images = get_images(Path(args.source))  # 获取所有待处理的图像文件
    print(f"[INFO] Found {len(images)} images")  # 打印找到的图像数量

    # =========================
    # inference loop
    # =========================
    for img_path in images:  # 遍历每个图像文件
        img = cv2.imread(str(img_path))  # 读取图像（BGR格式）
        if img is None:  # 如果图像读取失败
            continue  # 跳过该图像

        h, w = img.shape[:2]  # 获取原始图像的高度和宽度

        # preprocess（预处理）
        img_lb, ratio, pad = letterbox(img, 640)  # 对图像进行letterbox处理，缩放到640x640
        img_rgb = cv2.cvtColor(img_lb, cv2.COLOR_BGR2RGB)  # 将BGR格式转换为RGB格式

        inp = img_rgb.astype(np.float32) / 255.0  # 归一化到[0,1]范围并转换为float32类型
        inp = np.transpose(inp, (2, 0, 1))  # 调整维度顺序：从(H,W,C)转换为(C,H,W)
        inp = np.expand_dims(inp, 0)  # 添加批次维度，变为(1,C,H,W)
        inp = np.ascontiguousarray(inp)  # 确保数组在内存中是连续的

        # inference（推理）
        out = session.infer([inp])[0]  # 执行推理，获取输出（假设第一个输出是检测结果）
        pred = out[0]   # 获取预测结果，形状为(300,6)

        # postprocess（后处理）
        boxes, scores, cls_ids = postprocess(
            pred,  # 预测结果
            (h, w),  # 原始图像尺寸
            ratio,  # 缩放比例
            pad,  # 填充值
            args.conf,  # 置信度阈值
            args.iou  # IoU阈值
        )

        # =========================
        # save YOLO txt（保存YOLO格式标签）
        # =========================
        txt_path = label_dir / f"{img_path.stem}.txt"  # 生成标签文件路径（与图像同名，扩展名为.txt）
        lines = []  # 存储标签行的列表

        for box, score, cid in zip(boxes, scores, cls_ids):  # 遍历每个检测结果
            x1, y1, x2, y2 = box  # 获取边界框坐标

            # 转换为YOLO格式：中心点坐标和宽高（均归一化到[0,1]）
            cx = (x1 + x2) / 2 / w  # 中心点x坐标（归一化）
            cy = (y1 + y2) / 2 / h  # 中心点y坐标（归一化）
            bw = (x2 - x1) / w  # 边界框宽度（归一化）
            bh = (y2 - y1) / h  # 边界框高度（归一化）

            # 格式：类别ID 中心x 中心y 宽度 高度 置信度
            lines.append(
                f"{int(cid)} {cx:.6f} {cy:.6f} {bw:.6f} {bh:.6f} {score:.6f}"
            )

        txt_path.write_text("\n".join(lines))  # 将标签写入文件（每行一个检测结果）

        # =========================
        # save image（保存绘制了检测框的图像）
        # =========================
        vis = draw_boxes(img.copy(), boxes, scores, cls_ids)  # 在图像副本上绘制检测结果

        save_img_path = img_dir / f"{img_path.stem}.jpg"  # 生成保存图像路径（转换为jpg格式）
        cv2.imwrite(str(save_img_path), vis)  # 保存绘制后的图像

        print(f"{img_path.name}: {len(boxes)} objects")  # 打印检测到的目标数量


if __name__ == "__main__":  # 如果作为主程序运行
    main()  # 调用主函数