OAK相机如何将yoloV6模型转换成blob格式？（2.0 及之后版本）

编辑：OAK中国
首发：oakchina.cn
喜欢的话，请多多👍⭐️✍
内容可能会不定期更新，官网内容都是最新的，请查看首发地址链接。

▌前言

Hello，大家好，这里是OAK中国，我是助手君。

最近咱社群里有几个朋友在将yolox转换成blob的过程有点不清楚，所以我就写了这篇博客。（请夸我贴心！咱的原则：合理要求，有求必应！）

1.其他Yolo转换及使用教程请参考
2.检测类的yolo模型建议使用在线转换（地址），如果在线转换不成功，你再根据本教程来做本地转换。

▌`.pt` 转换为 `.onnx`

使用下列脚本(将脚本放到 YOLOv6 根目录中)将 pytorch 模型转换为 onnx 模型，若已安装 openvino_dev，则可进一步转换为 OpenVINO 模型：

示例用法：

python export_onnx.py -w <path_to_model>.pt -imgsz 640

export_onnx.py :

#!/usr/bin/env python3
# -*- coding:utf-8 -*-
import argparse
import json
import subprocess
import sys
import time
from pathlib import Path

import onnx
import torch
import torch.nn as nn

ROOT = Path.cwd()
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))

from yolov6.layers.common import *
from yolov6.models.efficientrep import (CSPBepBackbone, CSPBepBackbone_P6,
                                        EfficientRep, EfficientRep6)
from yolov6.models.effidehead import Detect
from yolov6.models.yolo import *
from yolov6.utils.checkpoint import load_checkpoint
from yolov6.utils.events import LOGGER


class YoloV6BackBone(nn.Module):
    """
    Backbone of YoloV6 model, it takes the model's original backbone and wraps it in this
    universal class. This was created for backwards compatibility with R2 models.
    """

    def __init__(self, old_layer, uses_fuse_P2=True, uses_6_erblock=False):
        super().__init__()

        self.uses_fuse_P2 = uses_fuse_P2
        self.uses_6_erblock = uses_6_erblock

        self.fuse_P2 = old_layer.fuse_P2 if hasattr(old_layer, "fuse_P2") else False

        self.stem = old_layer.stem
        self.ERBlock_2 = old_layer.ERBlock_2
        self.ERBlock_3 = old_layer.ERBlock_3
        self.ERBlock_4 = old_layer.ERBlock_4
        self.ERBlock_5 = old_layer.ERBlock_5
        if uses_6_erblock:
            self.ERBlock_6 = old_layer.ERBlock_6

    def forward(self, x):
        outputs = []
        x = self.stem(x)
        x = self.ERBlock_2(x)
        if self.uses_fuse_P2 and self.fuse_P2:
            outputs.append(x)
        elif not self.uses_fuse_P2:
            outputs.append(x)
        x = self.ERBlock_3(x)
        outputs.append(x)
        x = self.ERBlock_4(x)
        outputs.append(x)
        x = self.ERBlock_5(x)
        outputs.append(x)
        if self.uses_6_erblock:
            x = self.ERBlock_6(x)
            outputs.append(x)

        return tuple(outputs)


class DetectV6R2(nn.Module):
    """Efficient Decoupled Head for YOLOv6 R2&R3
    With hardware-aware degisn, the decoupled head is optimized with
    hybridchannels methods.
    """

    # def __init__(self, num_classes=80, anchors=1, num_layers=3, inplace=True, head_layers=None, use_dfl=True, reg_max=16):  # detection layer
    def __init__(self, old_detect):  # detection layer
        super().__init__()
        self.nc = old_detect.nc  # number of classes
        self.no = old_detect.no  # number of outputs per anchor
        self.nl = old_detect.nl  # number of detection layers
        if hasattr(old_detect, "anchors"):
            self.anchors = old_detect.anchors
        self.grid = old_detect.grid  # [torch.zeros(1)] * self.nl
        self.prior_prob = 1e-2
        self.inplace = old_detect.inplace
        stride = [8, 16, 32]  # strides computed during build
        self.stride = torch.tensor(stride)
        self.use_dfl = old_detect.use_dfl
        self.reg_max = old_detect.reg_max
        self.proj_conv = old_detect.proj_conv
        self.grid_cell_offset = 0.5
        self.grid_cell_size = 5.0

        # Init decouple head
        self.stems = old_detect.stems
        self.cls_convs = old_detect.cls_convs
        self.reg_convs = old_detect.reg_convs
        self.cls_preds = old_detect.cls_preds
        self.reg_preds = old_detect.reg_preds

    def forward(self, x):
        outputs = []
        for i in range(self.nl):
            b, _, h, w = x[i].shape
            l = h * w
            x[i] = self.stems[i](x[i])
            cls_x = x[i]
            reg_x = x[i]
            cls_feat = self.cls_convs[i](cls_x)
            cls_output = self.cls_preds[i](cls_feat)
            reg_feat = self.reg_convs[i](reg_x)
            reg_output = self.reg_preds[i](reg_feat)

            if self.use_dfl:
                reg_output = reg_output.reshape([-1, 4, self.reg_max + 1, l]).permute(
                    0, 2, 1, 3
                )
                reg_output = self.proj_conv(F.softmax(reg_output, dim=1))[:, 0]
                reg_output = reg_output.reshape([-1, 4, h, w])

            cls_output = torch.sigmoid(cls_output)
            conf, _ = cls_output.max(1, keepdim=True)
            output = torch.cat([reg_output, conf, cls_output], axis=1)
            outputs.append(output)

        return outputs


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument(
        "-w", "--weights", type=Path, default="./yolov6s.pt", help="weights path"
    )
    parser.add_argument(
        "-imgsz",
        "--img-size",
        nargs="+",
        type=int,
        default=[640, 640],
        help="image size",
    )  # height, width
    parser.add_argument(
        "--inplace", action="store_true", help="set Detect() inplace=True"
    )
    parser.add_argument("--opset", type=int, default=12, help="opset version")
    args = parser.parse_args()
    args.img_size *= 2 if len(args.img_size) == 1 else 1  # expand
    print(args)
    t = time.time()

    # Check device
    device = torch.device("cpu")
    # Load PyTorch model
    model = load_checkpoint(
        str(args.weights), map_location=device, inplace=True, fuse=True
    )  # load FP32 model

    labels = model.names  # get class names

    for layer in model.modules():
        if isinstance(layer, RepVGGBlock):
            layer.switch_to_deploy()

    for n, module in model.named_children():
        if isinstance(module, EfficientRep) or isinstance(module, CSPBepBackbone):
            setattr(model, n, YoloV6BackBone(module))
        elif isinstance(module, EfficientRep6):
            setattr(model, n, YoloV6BackBone(module, uses_6_erblock=True))
        elif isinstance(module, CSPBepBackbone_P6):
            setattr(
                model,
                n,
                YoloV6BackBone(module, uses_fuse_P2=False, uses_6_erblock=True),
            )

    # Input
    img = torch.zeros(1, 3, *args.img_size).to(
        device
    )  # image size(1,3,320,192) iDetection

    # Update model
    model.eval()
    for k, m in model.named_modules():
        if isinstance(m, Conv):  # assign export-friendly activations
            if isinstance(m.act, nn.SiLU):
                m.act = SiLU()
        elif isinstance(m, Detect):
            m.inplace = args.inplace

    model.detect = DetectV6R2(model.detect)
    num_branches = len(model.detect.grid)

    y = model(img)  # dry run

    # ONNX export
    try:
        LOGGER.info("\nStarting to export ONNX...")
        output_list = [f"output{i+1}_yolov6r2" for i in range(num_branches)]
        export_file = args.weights.with_suffix(".onnx")  # filename
        torch.onnx.export(
            model,
            img,
            export_file,
            verbose=False,
            opset_version=args.opset,
            training=torch.onnx.TrainingMode.EVAL,
            do_constant_folding=True,
            input_names=["images"],
            output_names=output_list,
            dynamic_axes=None,
        )

        # Checks
        onnx_model = onnx.load(export_file)  # load onnx model
        onnx.checker.check_model(onnx_model)  # check onnx model
        try:
            import onnxsim

            LOGGER.info("\nStarting to simplify ONNX...")
            onnx_model, check = onnxsim.simplify(onnx_model)
            assert check, "assert check failed"
        except Exception as e:
            LOGGER.info(f"Simplifier failure: {e}")
        LOGGER.info(f"ONNX export success, saved as {export_file}")
    except Exception as e:
        LOGGER.info(f"ONNX export failure: {e}")

    export_json = export_file.with_suffix(".json")
    export_json.with_suffix(".json").write_text(
        json.dumps(
            {
                "anchors": [],
                "anchor_masks": {},
                "coordinates": 4,
                "labels": labels,
                "num_classes": model.nc,
            },
            indent=4,
        )
    )
    LOGGER.info("Labels data export success, saved as %s" % export_json)

    # OpenVINO export
    LOGGER.info("\nStarting to export OpenVINO...")
    export_dir = Path(str(export_file).replace(".onnx", "_openvino"))
    OpenVINO_cmd = (
        "mo --input_model %s --output_dir %s --data_type FP16 --scale 255 --reverse_input_channel --output '%s' "
        % (export_file, export_dir, ",".join(output_list))
    )
    try:
        subprocess.check_output(OpenVINO_cmd, shell=True)
        LOGGER.info(f"OpenVINO export success, saved as {export_dir}")
    except Exception as e:
        LOGGER.info(f"OpenVINO export failure: {e}")
        LOGGER.info("\nBy the way, you can try to export OpenVINO use:")
        LOGGER.info("\n%s" % OpenVINO_cmd)

    # OAK Blob export
    LOGGER.info("\nThen you can try to export blob use:")
    export_xml = export_dir / export_file.with_suffix(".xml")
    export_blob = export_dir / export_file.with_suffix(".blob")
    blob_cmd = (
        "compile_tool -m %s -ip U8 -d MYRIAD -VPU_NUMBER_OF_SHAVES 6 -VPU_NUMBER_OF_CMX_SLICES 6 -o %s"
        % (export_xml, export_blob)
    )
    LOGGER.info("\n%s" % blob_cmd)

    # Finish
    LOGGER.info("\nExport complete (%.2fs)" % (time.time() - t))

可以使用 Netron 查看模型结构：

在这里插入图片描述

▌转换

openvino 本地转换

onnx -> openvino

mo 是 openvino_dev 2022.1 中脚本，

安装命令为 pip install openvino-dev

mo --input_model yolov6n.onnx --reverse_input_channel

openvino -> blob

<path>/compile_tool -m yolov6n.xml \
-ip U8 -d MYRIAD \
-VPU_NUMBER_OF_SHAVES 6 \
-VPU_NUMBER_OF_CMX_SLICES 6

在线转换

blobconvert 网页 http://blobconverter.luxonis.com/

进入网页，按下图指示操作：

在这里插入图片描述

修改参数，转换模型：

在这里插入图片描述

选择 onnx 模型
修改 optimizer_params 为 --data_type=FP16 --scale 255 --reverse_input_channel
修改 shaves 为 6
转换

blobconverter python 代码

blobconverter.from_onnx(
            "yolov6n.onnx",	
            optimizer_params=[
                "--scale 255",
                "--reverse_input_channel",
            ],
            shaves=6,
        )

blobconvert cli

	blobconverter --onnx yolov6n.onnx -sh 6 -o . --optimizer-params "scale=255 --reverse_input_channel"

▌DepthAI 示例

正确解码需要可配置的网络相关参数：

setNumClasses - YOLO 检测类别的数量
setIouThreshold - iou 阈值
setConfidenceThreshold - 置信度阈值，低于该阈值的对象将被过滤掉

import cv2
import depthai as dai
import numpy as np

model = dai.OpenVINO.Blob("yolov6n.blob")
dim = model.networkInputs.get("images").dims
W, H = dim[:2]
labelMap = [
    # "class_1","class_2","..."
    "class_%s"%i for i in range(80)
]

# Create pipeline
pipeline = dai.Pipeline()

# Define sources and outputs
camRgb = pipeline.create(dai.node.ColorCamera)
detectionNetwork = pipeline.create(dai.node.YoloDetectionNetwork)
xoutRgb = pipeline.create(dai.node.XLinkOut)
nnOut = pipeline.create(dai.node.XLinkOut)

xoutRgb.setStreamName("rgb")
nnOut.setStreamName("nn")

# Properties
camRgb.setPreviewSize(W, H)
camRgb.setResolution(dai.ColorCameraProperties.SensorResolution.THE_1080_P)
camRgb.setInterleaved(False)
camRgb.setColorOrder(dai.ColorCameraProperties.ColorOrder.BGR)
camRgb.setFps(40)

# Network specific settings
detectionNetwork.setBlob(model)
detectionNetwork.setConfidenceThreshold(0.5)
detectionNetwork.setNumClasses(80)
detectionNetwork.setCoordinateSize(4)
detectionNetwork.setAnchors([])
detectionNetwork.setAnchorMasks({})
detectionNetwork.setIouThreshold(0.5)

# Linking
camRgb.preview.link(detectionNetwork.input)
camRgb.preview.link(xoutRgb.input)
detectionNetwork.out.link(nnOut.input)

# Connect to device and start pipeline
with dai.Device(pipeline) as device:
    # Output queues will be used to get the rgb frames and nn data from the outputs defined above
    qRgb = device.getOutputQueue(name="rgb", maxSize=4, blocking=False)
    qDet = device.getOutputQueue(name="nn", maxSize=4, blocking=False)

    frame = None
    detections = []
    color2 = (255, 255, 255)

    # nn data, being the bounding box locations, are in <0..1> range - they need to be normalized with frame width/height
    def frameNorm(frame, bbox):
        normVals = np.full(len(bbox), frame.shape[0])
        normVals[::2] = frame.shape[1]
        return (np.clip(np.array(bbox), 0, 1) * normVals).astype(int)

    def displayFrame(name, frame):
        color = (255, 0, 0)
        for detection in detections:
            bbox = frameNorm(frame, (detection.xmin, detection.ymin, detection.xmax, detection.ymax))
            cv2.putText(frame, labelMap[detection.label], (bbox[0] + 10, bbox[1] + 20), cv2.FONT_HERSHEY_TRIPLEX, 0.5, 255)
            cv2.putText(frame, f"{int(detection.confidence * 100)}%", (bbox[0] + 10, bbox[1] + 40), cv2.FONT_HERSHEY_TRIPLEX, 0.5, 255)
            cv2.rectangle(frame, (bbox[0], bbox[1]), (bbox[2], bbox[3]), color, 2)
        # Show the frame
        cv2.imshow(name, frame)

    while True:
        inRgb = qRgb.tryGet()
        inDet = qDet.tryGet()

        if inRgb is not None:
            frame = inRgb.getCvFrame()

        if inDet is not None:
            detections = inDet.detections


        if frame is not None:
            displayFrame("rgb", frame)

        if cv2.waitKey(1) == ord('q'):
            break