+86-13951940532
contact@oakchina.cn

Yolov5 6.0 转换成blob格式

Yolov5 6.0 转换成blob格式

depthai api 需要的 yolo 模型 输出格式为 B, N*(x, y, h, w, box_score, class_no_1, …,class_no_N), Cx, Cy,其中:

  • B – 批量大小
  • N – 检测框数量
  • (x, y) – 检测框中心相对于单元格的坐标
  • h, w – 框的原始高度和宽度,然后乘以相应的锚点得到绝对高度和宽度值
  • box_score – 检测框置信度,在 [0,1] 范围内
  • class_no_1, …, class_no_N – 概率分布,在 [0,1] 范围内,乘以置信度得到每个类的置信度
  • Cx, Cy – 检测框索引

可以自行转换 onnx 或者使用 release 提供的onnx

python export.py --simplify --opset 12 --include onnx --batch-size 1 --imgsz 640 --weights yolov5n.pt

现在我们将阅读和编辑我们的 ONNX 模型,但首先让我们谈谈 YoloV5 模型的结构。与其他的 YoloV3/YoloV4 不同,YoloV5 的实现已经包含了计算边界框的所有必要步骤。因此,您所要做的就是在输出上使用 NMS 并计算检测到的对象的最终框。

然而,DepthAI 中的 YoloDetectionNetwork 节点最初是为 YoloV3 和 V4 开发的。它从最后一个卷积层获取输出,并使用掩码和锚点以及设备上的 NMS 进行所有必要的计算。 YoloV3/V4 中的 box 计算也与 YoloV5 中的计算略有不同,因为后者使用乘法和幂运算以避免之前版本中可能出现的 exp 运算问题。盒子中心坐标的方程也略有不同。因此,我们必须稍微编辑 YoloV5 模型的输出。

如果您研究 ONNX 模型的架构(我们在 Netron 中这样做),您会发现 3 个卷积层略高于输出。这些是我们感兴趣且形状正确的层,它们下面的层用于后期处理,我们不需要。但是,请注意在此后处理步骤之前的 Sigmoid 层,但在 reshape 操作之后。理想情况下,我们会在卷积层之后使用 Sigmoid。

我们对最后三层感兴趣,它们代表我们新的 sigmoid 层的输入。

import onnx

onnx_model = onnx.load("yolov5n.onnx")

conv_indices = []
for i, n in enumerate(onnx_model.graph.node):
  if "Conv" in n.name:
    conv_indices.append(i)

input1, input2, input3 = conv_indices[-3:]

sigmoid1 = onnx.helper.make_node(
    'Sigmoid',
    inputs=[onnx_model.graph.node[input1].output[0]],
    outputs=['output1_yolov5'],
)

sigmoid2 = onnx.helper.make_node(
    'Sigmoid',
    inputs=[onnx_model.graph.node[input2].output[0]],
    outputs=['output2_yolov5'],
)

sigmoid3 = onnx.helper.make_node(
    'Sigmoid',
    inputs=[onnx_model.graph.node[input3].output[0]],
    outputs=['output3_yolov5'],
)

onnx_model.graph.node.append(sigmoid1)
onnx_model.graph.node.append(sigmoid2)
onnx_model.graph.node.append(sigmoid3)

onnx.save(onnx_model, "yolov5n.onnx")

首先,我们加载我们之前导出的模型并收集所有卷积层的索引。

我们创建了 3 个新的 sigmoid 层,并将它们连接到带有输入标志的卷积层。 我们还重命名了输出,以便它们包含“_yolov5”,DepthAI 使用它来区分 YoloV3/V4 和 YoloV5。 这将告诉设备应用正确的操作(乘法、幂、…而不是 exp)并返回正确的预测。 我们将这些节点附加到模型中,并将编辑后的 ONNX 文件保存到“onnx_output_path”(从技术上讲,它还没有被修剪,但将会是😀)。

如果我们再次研究 Netron 中的模型,我们可以看到添加的带有集合名称的 sigmoid 层。

剩下要做的就是将模型导出到 OpenVINO IR 并生成我们可以与 DepthAI API 一起使用的 blob! 💻

YOLOv5 6.0

openvino 本地转换

onnx -> IR

mo.py --input_model yolov5n.onnx --scale 255 --reverse_input_channel --output "output1_yolov5,output2_yolov5,output3_yolov5"

IR -> Blob

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

blobconvert 代码 onnx -> Blob

blobconverter.from_onnx(
            "yolov5n.onnx",
            optimizer_params=[
                "--scale=255",
                "--reverse_input_channel",
                "--output=output1_yolov5,output2_yolov5,output3_yolov5",
            ],
            data_type="FP32",
            shaves=6,
        )

blobconvert Cli

blobconverter --onnx yolov5n.onnx \
-dt FP32 -sh 6 -o . \
--optimizer-params "scale=255 --reverse_input_channel --output=output1_yolov5,output2_yolov5,output3_yolov5"

DepthAI 设置

import depthai as dai

pipeline = dai.Pipeline()
yolo = pipeline.create(dai.node.YoloDetectionNetwork)
yolo.setBlobPath("yolov5n.blob")
yolo.setConfidenceThreshold(0.5)
yolo.setNumClasses(80)
yolo.setCoordinateSize(4)
yolo.setAnchors([10,13, 16,30, 33,23, 30,61, 62,45, 59,119, 116,90, 156,198, 373,326])
yolo.setAnchorMasks({
        		"side80" : [0,1,2],
                "side40" : [3,4,5],
                "side20" : [6,7,8]
                })
yolo.setIouThreshold(0.5)
  • input_size – 模型的输入形状,
  • num_classes – Yolo 可以检测到的类数,
  • coordinates – 坐标的个数(默认为 4),
  • anchors – yolo 锚点(位于 CFG 中的 [yolo] 层),
  • iou_threshold – iou 置信阈值,
  • confidence_threshold – 检测到对象的置信阈值,
  • anchor_masks – 锚掩码.

注意: 参数值必须与训练期间在 CFG 中设置的值相匹配。如果使用不同的输入宽度,还应该将side32 改为 sideX ,将 side16 改为 sideY, 其中 X = width/16 , Y = width/32. 。如果您使用的是非微型模型,那么这些值是width/8, width/16, 和 width/32.

Tags: