YOLOv5x6模型来了! 同样支持CPU上ONNX部署与推理

最新推荐文章于 2024-07-25 16:59:12 发布

3Ｄ视觉工坊

最新推荐文章于 2024-07-25 16:59:12 发布

阅读量2.5k

点赞数

文章标签：人工智能编程语言 etag qml iot

点击上方“3D视觉工坊”，选择“星标”
干货第一时间送达

说一下是YOLOv5的第五个版本，不是YOLO的第五个版本！是YOLOv5又又改进了！

YOLOv5x6模型来了

自从Pytorch版本YOLOv5发布之后，经历过了四个版本的升级，YOLOv5的功能与模型精度不断提升。不久之前YOLOv5-Pytorch发布第五个版本，第五个版本跟之前版本最大的差异就是多出了一个输出层，之前的输出层分辨率倍数为：[8、16、32]三个层的输出。现在YOLOv5多出的一个输出层之后，它的输出变为：[8、16、32、64]。但是我找遍了model文件夹中的模型描述文件，没发现跟模型对应的描述文件，所以我就不知道是否可以支持训练。这里我就重点说一下如何使用已经发布的YOLOv5支持四层输出的模型：

别担心之前的第四版本模型都可以用，模型命名带数字6的都是支持四个输出层的，以YOLOv5s6为例，转ONNX格式查看截图如下：

数值精度取舍

首先需要需要运行下面的命令行获取整个YOLOv5项目的源码：

git clone https://github.com/ultralytics/yolov5.git

然后测试运行：

测试没有问题，就可以转换模型yolov5s6为ONNX格式，命令行如下：

这样就获取到了ONNX格式文件，注意它的对应输出格式为：

1x3x1280x1280RGB, 像素值在0到1之间

导入模型

model_onnx = "D:/python/yolov5/yolov5s6.onnx"net = ie.read_network(model=model_onnx)input_blob = next(iter(net.input_info))head_it = iter(net.outputs)out_blob_160 = next(head_it) # 160out_blob_80 = next(head_it)  # 80out_blob_40 = next(head_it)  # 40out_blob_20 = next(head_it)  # 20

处理输入图象与推理

rgb = cv.cvtColor(frame, cv.COLOR_BGR2RGB)rgb = cv.resize(rgb, (w, h))image = np.float32(rgb.transpose(2, 0, 1))/255.0inf_start = time.time()res = exec_net.infer(inputs={input_blob:[image]})outs = []outs.append(res[out_blob_20])outs.append(res[out_blob_40])outs.append(res[out_blob_80])outs.append(res[out_blob_160])

解析输出结果：

for out in outs:
    out = np.squeeze(out, 0)
    dims, side_h, side_w, side_data = out.shape
    side_square = side_h * side_w
    stride = get_stride(side_w)
    anchor_index = get_anchor_index(side_w)
    for i in range(side_square):
        for d in range(dims):
            row = i // side_h;
            col = i % side_h;
            box_data = out[d, row, col]
            conf = sigmoid(box_data[4]);
            if conf < 0.25:
                continue
            xx = (sigmoid(box_data[0]) * 2 - 0.5 + col) * stride;
            yy = (sigmoid(box_data[1]) * 2 - 0.5 + row) * stride;
            ww = pow(sigmoid(box_data[2]) * 2, 2) * anchors[anchor_index + d * 2];
            hh = pow(sigmoid(box_data[3]) * 2, 2) * anchors[anchor_index + d * 2 + 1];

            max_prob = -1
            max_index = -1
            for p in range(5, 85, 1):
                prob = sigmoid(box_data[p])
                if prob > max_prob:
                    max_prob = prob
                    max_index = p - 5

            # 转换为top - left, bottom - right坐标
            x1 = int((xx - ww / 2) * sx)
            y1 = int((yy - hh / 2) * sy)
            x2 = int((xx + ww / 2) * sx)
            y2 = int((yy + hh / 2) * sy)

            boxes.append([x1, y1, x2-x1, y2-y1])
            confidences.append(float(conf))
            classIds.append(max_index)

最后NMS输出结果如下：