目标检测 YOLOv5 - 最新版本v6.2模型在瑞芯微 Rockchip设备上运行的方案

目标检测 YOLOv5 - 最新版本v6.2模型在瑞芯微 Rockchip设备上运行的方案

包含pytorch模型转rknn

flyfish

代码

https://github.com/shaoshengsong/rockchip_rknn_yolov5

模型的导出

本例是以640x640的输入来说明

git clone https://github.com/ultralytics/yolov5  # clone
cd yolov5
pip install -r requirements.txt  # install

可以切换到v6.2然后执行导出onnx

python export.py --weights yolov5s.pt --include torchscript onnx

或者直接下载v6.2的源码

需要更改导出部分
这个为了说明如何编写代码，导出了4部分。在使用时要么输出3部分，要么输出1部分，只选其一。

方案1 模型的输出如下

[1,25200,85] 
'
运行
运行

主要是后处理的编写

def yolov5_post_process(input_image, outputs):
      class_ids = []
      confidences = []
      boxes = []
      # Rows.
      rows = outputs[0].shape[1]

      # Iterate through detections.
      for r in range(rows):
            row = outputs[0][0][r]
            confidence = row[4]
            # Discard bad detections and continue.
            if confidence >= CONFIDENCE_THRESHOLD:
                  classes_scores = row[5:]
                  # Get the index of max class score.
                  class_id = np.argmax(classes_scores)
                  #  Continue if the class score is above threshold.
                  if (classes_scores[class_id] > SCORE_THRESHOLD):
                        confidences.append(confidence * classes_scores[class_id])
                        class_ids.append(class_id)
                        cx, cy, w, h = row[0], row[1], row[2], row[3]
                        left = int((cx - w/2) * x_factor)
                        top = int((cy - h/2) * y_factor)
                        width = int(w * x_factor)
                        height = int(h * y_factor)
                        box = np.array([left, top, width, height])
                        boxes.append(box)

      indices = cv2.dnn.NMSBoxes(boxes, confidences, CONFIDENCE_THRESHOLD, NMS_THRESHOLD)
      for i in indices:
            box = boxes[i]
            left = box[0]
            top = box[1]
            width = box[2]
            height = box[3]             
                   
            cv2.rectangle(input_image, (left, top), (left + width, top + height), (255, 0, 0), 2)
                              
            label = "{}:{:.2f}".format(CLASSES[class_ids[i]], confidences[i])             

            cv2.putText(input_image, label,
                    (left,top),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    0.6, (0, 0, 255), 2)
      return input_image
'
运行
运行

方案1完整的pytorch模型转rknn代码下载地址

https://github.com/shaoshengsong/rockchip_rknn_yolov5 

方案2 模型的输出如下

 (1, 3, 80, 80, 85)
 (1, 3, 40, 40, 85)
 (1, 3, 20, 20, 85)

在写代码时根据输出的name编写

ret = rknn.load_onnx(model=ONNX_MODEL, outputs=['339', '377', '415'])
if ret != 0:
    print('Load model failed!')
    exit(ret)
print('done')

中间转换的shape

(3, 80, 80, 85)
(3, 40, 40, 85)
(3, 20, 20, 85)
'
运行
运行

最终使用的shape

(80, 80, 3, 85)
(40, 40, 3, 85)
(20, 20, 3, 85)
'
运行
运行

方案2完整的pytorch模型转rknn代码下载地址

https://github.com/shaoshengsong/rockchip_rknn_yolov5 

在yolov5_v6.2文件夹中

方案3

需要在Ubuntu下进行预编译和量化
为了Rockchip的设备使用量化，这里把box和score分开了。
box范围是0-640
score的范围是0-1
box 由[x_center, y_center, width, height] 变成了[left, top, right, bottom]

常见的边框（bounding box ）坐标表示方法

最后模型导出的样子如下，输出两部分

pytorch原模型推理结果

rknn推理结果