本文训练的事yolov5-v7实例分割模型的详细记录，
关于yolov5-v7的官方代码可从官方githu下载
https://github.com/ultralytics/yolov5/tree/v7.0
由于我的实际场景需要部署在rk3588开发板上，所以使用的是瑞芯微官方github的代码
# 克隆master分支
https://github.com/airockchip/yolov5?tab=readme-ov-file
说明：瑞芯微官方的yolov5代码是在yolov5官方（v7.0）代码基础上修改而来，如果想要部署在rk3588等系列开发板上的话，使用瑞芯微官方提供的代码可以直接训练然后导出，不用再进行相关修改，而使用yolo官方提供的代码需要进行相关的修改，两个项目训练流程并无差异。
# 瑞芯微官方相关的model zoo
https://github.com/airockchip/rknn_model_zoo/tree/main/examples/yolov5_seg
数据集标注
制作分割数据集需要利用到labelme软件，我用的版本是3.16.7
虚拟环境中运行下属指令：
# ubutnu
pip install pycocotools
# Windows
pip install pycocotools-windows
# 安装3.16.7
pip install labelme==3.16.7
运行labelme
labelme
得到的json文件格式如下所示：
将标注文件有.json格式转换为.txt格式
import json
import glob
import os
import cv2
import numpy as np
json_path = r"/home/data/project/customer_AAA/rk3588/yolov5-airockchip/Data_save/data_nut_bolt"; #此处填写存放json文件的地址
txt_dir=r"/home/data/project/customer_AAA/rk3588/yolov5-airockchip/Data_save/txts" # 用于存放生成的txt文件的目录
labels = ['bolt','nut']#此处填写你标注的标签名称 除了修改这里，还需要修改下面的for循环中类别相关内容
json_files = glob.glob(json_path + "/*.json")
for json_file in json_files:
print(json_file)
f = open(json_file)
json_info = json.load(f)
# print(json_info.keys())
img = cv2.imread(os.path.join(json_path, json_info["imagePath"]))
height, width, _ = img.shape
np_w_h = np.array([[width, height]], np.int32)
file_name = os.path.basename(json_file)
txt_name = file_name.replace(".json", ".txt")
txt_path=os.path.join(txt_dir,txt_name)
f = open(txt_path, "w")
txt_content = ""
for point_json in json_info["shapes"]:
np_points = np.array(point_json["points"], np.int32)
norm_points = np_points / np_w_h
norm_points_list = norm_points.tolist()
if point_json['label'] == labels[0]:
txt_content += "0 " + " ".join([" ".join([str(cell[0]), str(cell[1])]) for cell in norm_points_list]) + "\n"
elif point_json['label'] == labels[1]:
txt_content += "1 " + " ".join([" ".join([str(cell[0]), str(cell[1])]) for cell in norm_points_list]) + "\n"
f.write(txt_content)
转换完成后，会得到对应于每个图像的**.txt**文件
每一行的数据含义如下所示：
类别标签 归一化后的点坐标
数据集划分：
执行如下代码：
import os
import random
import shutil
rootpath = r'/home/data/project/customer_AAA/rk3588/yolov5-airockchip/Data_save/data_nut_bolt_result/'#此处为img和txt文件夹存放位置，地址后面要有/结尾
set1 = ['images','labels']
set2 = ['train','val']
for s1 in set1:
if not os.path.exists(rootpath+s1):
os.mkdir(rootpath+s1)
for s2 in set2:
if not os.path.exists(rootpath+s1+'/'+s2):
os.mkdir(rootpath+s1+'/'+s2)
# 这是原始图片路径
img_path = rootpath+'img'
# 这是生成的txt路径
txt_path = rootpath+'txt'
file_names = os.listdir(img_path)
l = 0.8
n = len(file_names)
train_files = random.sample(file_names, int(n*l))
for file in file_names:
print(file)
if not os.path.exists(txt_path+'/'+file[:-3]+'txt'):
os.remove(img_path+'/'+file)
print(file[:-3]+'txt,不存在')
continue
if file in train_files:
shutil.copy(img_path+'/'+file,rootpath+'images/train/'+file)
shutil.copy(txt_path+'/'+file[:-3]+'txt',rootpath+'labels/train/'+file[:-3]+'txt')
else:
shutil.copy(img_path+'/'+file,rootpath+'images/val/'+file)
shutil.copy(txt_path+'/'+file[:-3]+'txt',rootpath+'labels/val/'+file[:-3]+'txt')
print('ok!!')
print(len(train_files))
上述代码中，需要注意rootpath的填写路径，因为路径用到的地方是字符串拼接
执行完毕后会在同级目录下生成对应的目录文件
设置数据集配置文件
在目录yolov5/data下
将coco128-seg.yaml配置文件备份，命名为coco128-seg_nut_bolt.yaml
# YOLOv5 🚀 by Ultralytics, GPL-3.0 license
# COCO128-seg dataset https://www.kaggle.com/ultralytics/coco128 (first 128 images from COCO train2017) by Ultralytics
# Example usage: python train.py --data coco128.yaml
# parent
# ├── yolov5
# └── datasets
# └── coco128-seg ← downloads here (7 MB)
# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
path: /home/data/project/customer_AAA/rk3588/yolov5-airockchip/Data_seg/data_nut_bolt # dataset root dir
train: images/train # train images (relative to 'path') 128 images
val: images/val # val images (relative to 'path') 128 images
test: # test images (optional)
# Classes
names:
0: bolt
1: nut
由于我的项目是进行螺丝、螺母分割，所以填写的识别类别是两类。
设置模型配置文件
在目录yolov5/models下，拷贝yolov5m-seg.yaml文件重命名为yolov5m-seg.yaml
(python3.8-tk2-2.0) root@f95e42ca3a90:/home/data/project/customer_AAA/rk3588/yolov5-airockchip/models# tree
.
|-- __init__.py
|-- __pycache__
| |-- __init__.cpython-38.pyc
| |-- common.cpython-38.pyc
| |-- common_rk_plug_in.cpython-38.pyc
| |-- experimental.cpython-38.pyc
| `-- yolo.cpython-38.pyc
|-- common.py
|-- common_rk_plug_in.py
|-- experimental.py
|-- hub
| |-- anchors.yaml
| |-- yolov3-spp.yaml
| |-- yolov3-tiny.yaml
| |-- yolov3.yaml
| |-- yolov5-bifpn.yaml
| |-- yolov5-fpn.yaml
| |-- yolov5-p2.yaml
| |-- yolov5-p34.yaml
| |-- yolov5-p6.yaml
| |-- yolov5-p7.yaml
| |-- yolov5-panet.yaml
| |-- yolov5l6.yaml
| |-- yolov5m6.yaml
| |-- yolov5n6.yaml
| |-- yolov5s-LeakyReLU.yaml
| |-- yolov5s-ghost.yaml
| |-- yolov5s-transformer.yaml
| |-- yolov5s6.yaml
| `-- yolov5x6.yaml
|-- segment
| |-- yolov5l-seg.yaml
| |-- yolov5m-seg.yaml
| |-- yolov5n-seg.yaml
| |-- yolov5n-seg_nut_bolt.yaml
| |-- yolov5s-seg.yaml
| `-- yolov5x-seg.yaml
|-- tf.py
|-- yolo.py
|-- yolov5l.yaml
|-- yolov5m.yaml
|-- yolov5n.yaml
|-- yolov5s.yaml
|-- yolov5s_2007_2012.yaml
`-- yolov5x.yaml
配置文件内容如下所示：
# YOLOv5 🚀 by Ultralytics, GPL-3.0 license
# Parameters
nc: 2 # number of classes
depth_multiple: 0.33 # model depth multiple
width_multiple: 0.25 # layer channel multiple
anchors:
- [10,13, 16,30, 33,23] # P3/8
- [30,61, 62,45, 59,119] # P4/16
- [116,90, 156,198, 373,326] # P5/32
# YOLOv5 v6.0 backbone
backbone:
# [from, number, module, args]
[[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
[-1, 1, Conv, [128, 3, 2]], # 1-P2/4
[-1, 3, C3, [128]],
[-1, 1, Conv, [256, 3, 2]], # 3-P3/8
[-1, 6, C3, [256]],
[-1, 1, Conv, [512, 3, 2]], # 5-P4/16
[-1, 9, C3, [512]],
[-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
[-1, 3, C3, [1024]],
[-1, 1, SPPF, [1024, 5]], # 9
]
# YOLOv5 v6.0 head
head:
[[-1, 1, Conv, [512, 1, 1]],
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
[[-1, 6], 1, Concat, [1]], # cat backbone P4
[-1, 3, C3, [512, False]], # 13
[-1, 1, Conv, [256, 1, 1]],
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
[[-1, 4], 1, Concat, [1]], # cat backbone P3
[-1, 3, C3, [256, False]], # 17 (P3/8-small)
[-1, 1, Conv, [256, 3, 2]],
[[-1, 14], 1, Concat, [1]], # cat head P4
[-1, 3, C3, [512, False]], # 20 (P4/16-medium)
[-1, 1, Conv, [512, 3, 2]],
[[-1, 10], 1, Concat, [1]], # cat head P5
[-1, 3, C3, [1024, False]], # 23 (P5/32-large)
[[17, 20, 23], 1, Segment, [nc, anchors, 32, 256]], # Detect(P3, P4, P5)
]
开启模型训练
在项目根目录下执行以下指令开启模型训练：
torchrun --nproc-per-node 2 ./segment/train.py --workers 0 --data data/coco128-seg_nut_bolt.yaml --cfg models/segment/yolov5n-seg_nut_bolt.yaml --img 640 --weights yolov5n-seg.pt --batch-size 128 --epochs 100
训练开始，绘制的图像如下所示：
训练的各个阶段曲线如下所示：
模型验证集统计数据如下所示：
Validating runs/train-seg/exp/weights/best.pt...
Fusing layers...
YOLOv5n-seg_nut_bolt summary: 224 layers, 1881103 parameters, 0 gradients, 6.7 GFLOPs
Class Images Instances Box(P R mAP50 mAP50-95) Mask(P R mAP50 mAP50-95): 100%|██████████| 1/1 00:03
all 85 604 0.995 0.998 0.994 0.841 0.972 0.975 0.962 0.736
bolt 85 300 0.994 0.997 0.995 0.856 0.947 0.95 0.93 0.641
nut 85 304 0.997 1 0.993 0.826 0.997 1 0.993 0.831
Results saved to runs/train-seg/exp
模型测试
执行如下检测指令
python segment/predict.py --weights runs/train-seg/exp/weights/best.pt --source Data_seg/data_nut_bolt/images/val/6.jpg
检测效果如下所示：
rk3588开发板端部署yolov5分割模型
将训练的模型转换成.rknn格式的模型文件
**注意：**上述模型是基于瑞芯微官方的yolov5（7.0）版本训练得到，所以直接可以进行转换成onnx格式，如果是yolo官方的yolov5代码，则需要改动相关的结构再转换
先转换成onnx格式，执行代码如下：
# 项目根目录下执行，得到onnx文件
python export.py --rknpu --weight runs/train-seg/exp/weights/yolov5n-seg.pt
再转换成.rknn格式文件
依托于瑞芯微github官方项目
https://github.com/airockchip/rknn_model_zoo/tree/v2.0.0
采用官方提供的yolov5-seg示例
其中yolov5-seg项目目录结构如下所示：
(python3.8-tk2-2.0) root@f95e42ca3a90:/home/data/project/customer_AAA/rk3588/rknn_model_zoo-2.0.0/examples/yolov5_seg# tree
.
|-- README.md
|-- cpp
| |-- CMakeLists.txt
| |-- easy_timer.h
| |-- main.cc
| |-- postprocess.h
| |-- rknpu1
| | |-- postprocess.cc
| | `-- yolov5_seg.cc
| |-- rknpu2
| | |-- postprocess.cc
| | `-- yolov5_seg.cc
| `-- yolov5_seg.h
|-- model
| |-- anchors_yolov5.txt
| |-- bus.jpg
| |-- coco_80_labels_list.txt
| `-- download_model.sh
|-- model_comparison
| |-- yolov5_seg_graph_comparison.jpg
| `-- yolov5_seg_output_comparison.jpg
|-- python
| |-- convert.py
| `-- yolov5_seg.py
`-- reference_results
|-- yolov5s_seg_c_demo_result.png
`-- yolov5s_seg_python_demo_result.png
7 directories, 20 files
将上述转换得到的yolov5n-seg.onnx文件拷贝至rknn_model_zoo-2.0.0/examples/yolov5_seg/python目录下：
python convert.py ../model/yolov5s-seg.onnx rk3588
对convert.py脚本进行修改，最终如下所示
import sys
from rknn.api import RKNN
DATASET_PATH = './data_list.txt' # 测试图像路径列表
DEFAULT_RKNN_PATH = './yolov5_seg.rknn'
DEFAULT_QUANT = True
def parse_arg():
if len(sys.argv) < 3:
print("Usage: python3 {} onnx_model_path [platform] [dtype(optional)] [output_rknn_path(optional)]".format(sys.argv[0]));
print(" platform choose from [rk3562, rk3566, rk3568, rk3588, rk1808, rv1109, rv1126]")
print(" dtype choose from [i8, fp] for [rk3562,rk3566,rk3568,rk3588]")
print(" dtype choose from [u8, fp] for [rk1808,rv1109,rv1126]")
exit(1)
model_path = sys.argv[1]
platform = sys.argv[2]
do_quant = DEFAULT_QUANT
if len(sys.argv) > 3:
model_type = sys.argv[3]
if model_type not in ['i8', 'u8', 'fp']:
print("ERROR: Invalid model type: {}".format(model_type))
exit(1)
elif model_type in ['i8', 'u8']:
do_quant = True
else:
do_quant = False
if len(sys.argv) > 4:
output_path = sys.argv[4]
else:
output_path = DEFAULT_RKNN_PATH
return model_path, platform, do_quant, output_path
if __name__ == '__main__':
model_path, platform, do_quant, output_path = parse_arg()
# Create RKNN object
rknn = RKNN(verbose=False)
# Pre-process config
print('--> Config model')
rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]], target_platform=platform)
print('done')
# Load model
print('--> Loading model')
ret = rknn.load_onnx(model=model_path)
if ret != 0:
print('Load model failed!')
exit(ret)
print('done')
# Build model
print('--> Building model')
ret = rknn.build(do_quantization=do_quant, dataset=DATASET_PATH)
if ret != 0:
print('Build model failed!')
exit(ret)
print('done')
# Export rknn model
print('--> Export rknn model')
ret = rknn.export_rknn(output_path)
if ret != 0:
print('Export rknn model failed!')
exit(ret)
print('done')
# Release
rknn.release()
执行如下指令：
python convert.py ./yolov5n-seg.onnx rk3588
测试脚本：
# 加载onnx文件测试
python yolov5_seg.py --model_path yolov5n-seg.onnx --img_folder data_set --img_save
分割效果一般，估计训练100轮次还是不够，哈哈哈
加载
python yolov5_seg.py --model_path yolov5_seg.rknn --img_folder /home/data/project/customer_AAA/rk3588/yolov5-airockchip/Data_save/data_nut_bolt_result/images/train --target rk3588 --img_save