训练
数据集最开始是下面的结构
datasets/
├── images/
├── labels/
注意如果数据集原始的是xml格式需要转换为txt格式,转换代码如下xml2yolotxt.py,会生成labels_txt目录和对应的标注文件。记得改名字labels_txt为labels
import os
import xml.etree.ElementTree as ET
def xml_to_yolo(xml_path, output_dir, class_names):
tree = ET.parse(xml_path)
root = tree.getroot()
# 获取图片尺寸
size = root.find('size')
img_w = int(size.find('width').text)
img_h = int(size.find('height').text)
# 创建YOLO格式的txt文件
txt_path = os.path.join(output_dir, os.path.splitext(os.path.basename(xml_path))[0] + '.txt')
with open(txt_path, 'w') as f:
for obj in root.iter('object'):
cls = obj.find('name').text
if cls not in class_names:
continue
cls_id = 0
# 获取边界框坐标
bbox = obj.find('bndbox')
xmin = int(bbox.find('xmin').text)
ymin = int(bbox.find('ymin').text)
xmax = int(bbox.find('xmax').text)
ymax = int(bbox.find('ymax').text)
# 转换为YOLO格式(归一化中心坐标和宽高)
x_center = ((xmin + xmax) / 2) / img_w
y_center = ((ymin + ymax) / 2) / img_h
width = (xmax - xmin) / img_w
height = (ymax - ymin) / img_h
f.write(f"{cls_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}\n")
# 使用示例
class_names = ["liner","container ship","bulk carrier","sailboat","other ship"] # 你的类别列表
xml_dir = "labels_xml"
output_dir = "labels_txt"
os.makedirs(output_dir, exist_ok=True)
for xml_file in os.listdir(xml_dir):
if xml_file.endswith('.xml'):
xml_to_yolo(os.path.join(xml_dir, xml_file), output_dir, class_names)
数据集比例划分
将数据集按照一定的比例分为训练数据集和验证数据集,通过下面的代码进行划分
import os
import shutil
from sklearn.model_selection import train_test_split
def split_dataset(images_dir, labels_dir, output_dir, train_ratio=0.8, random_seed=42):
"""
划分数据集为train/val,并保持images和labels的严格对应
参数:
images_dir: 原始图片路径 (如 "./images")
labels_dir: 原始标签路径 (如 "./labels")
output_dir: 输出根目录 (如 "./dataset")
train_ratio: 训练集比例 (默认0.8)
random_seed: 随机种子 (确保可复现)
"""
# 创建输出目录
os.makedirs(os.path.join(output_dir, "images/train"), exist_ok=True)
os.makedirs(os.path.join(output_dir, "images/val"), exist_ok=True)
os.makedirs(os.path.join(output_dir, "labels/train"), exist_ok=True)
os.makedirs(os.path.join(output_dir, "labels/val"), exist_ok=True)
# 获取所有图片文件名(不带扩展名)
image_files = [os.path.splitext(f)[0] for f in os.listdir(images_dir) if f.endswith(('.jpg', '.png', '.jpeg'))]
image_files.sort() # 确保顺序一致
# 划分train/val
train_files, val_files = train_test_split(
image_files,
train_size=train_ratio,
random_state=random_seed
)
# 复制文件函数
def copy_files(files, split_type):
for file in files:
# 复制图片
for ext in ['.jpg', '.png', '.jpeg']:
src_img = os.path.join(images_dir, file + ext)
if os.path.exists(src_img):
shutil.copy(src_img, os.path.join(output_dir, "images", split_type))
break
# 复制标签
src_label = os.path.join(labels_dir, file + ".txt")
if os.path.exists(src_label):
shutil.copy(src_label, os.path.join(output_dir, "labels", split_type))
# 执行复制
copy_files(train_files, "train")
copy_files(val_files, "val")
print(f"划分完成!\n训练集: {len(train_files)} 个样本\n验证集: {len(val_files)} 个样本")
if __name__ == "__main__":
# 使用示例
split_dataset(
images_dir="images", # 原始图片路径
labels_dir="labels", # 原始标签路径
output_dir="dataset", # 输出路径
train_ratio=0.8 # 训练集比例
)
然后在当前目录添加配置文件yolov8n.yaml和data.yaml内容分别为
# 从官方模型继承并修改类别数
nc: 1 # 与data.yaml中的类别数一致
# 数据集路径和类别定义
path: /workspace/to_docker/dataset_ship/dataset
train: images/train
val: images/val# 类别数量和名称
names:
0: ship
拷贝一个yolov8n.pt模型到当前目录,然后使用命令 yolo task=detect mode=train model=yolov8n.pt data=data.yaml epochs=100 imgsz=640进行训练
模型导出
使用 airockchip/ultralytics_yolov8 可以直接导出适配 rknpu 的模型,在 npu 上获得更高的推理效率。
该仓库对模型进行了优化:
• dfl 结构在 NPU 处理上性能不佳,移至模型外部。
• 假设有 6000 个候选框,原模型将 dfl 结构放置于’’框置信度过滤”前,则 6000 个候选框
都需要计算经过 dfl 计算;而将 dfl 结构放置于’’框置信度过滤”后,假设过滤后剩 100 个
候选框,则 dfl 部分计算量减少至 100 个,大幅减少了计算资源、带宽资源的占用。
将训练好的best.pt模型拷贝到ultralytics_yolov8目录
cd ultralytics_yolov8
export PYTHONPATH=./
#如果激活了rknn的环境需要通过deactivate命令取消激活
python ./ultralytics/engine/exporter.py
运行后会在 ultralytics_yolov8目录生成best_rknnopt.torchscript文件,将文件拷贝并改名到rk3588_yolov8/yolov8n_rknnopt.pt,然后运行命令
#pt2rknn.py
import sys
from rknn.api import RKNN
DATASET_PATH = '/home/ubuntu/3588_pt2rknn_yolov8/onnx2rknn/rknn_model_zoo-2.2.0/datasets/COCO/coco_subset_20.txt' #不用管 验证改false就行
DEFAULT_RKNN_PATH = './pt2rknn.rknn' #模型输出的路径
DEFAULT_QUANT = False #应该是用来验证的开关
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,rk3576,rk3588,rk1808,rv1109,rv1126]")
print(" dtype choose from [i8, fp] for [rk3562,rk3566,rk3568,rk3576,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')
# 加载Pytorch模型,如果是使用airockchip/ultralytics_yolov8的main分支,将导入onnx模型
print('--> Loading model')
#ret = rknn.load_onnx(model=model_path)
ret = rknn.load_pytorch(model=model_path, input_size_list=[[1, 3, 640, 640]])
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')
# 精度分析,,输出目录./snapshot
#print('--> Accuracy analysis')
#ret = rknn.accuracy_analysis(inputs=['./subset/000000052891.jpg'])
#if ret != 0:
# print('Accuracy analysis failed!')
# exit(ret)
#print('done')
# Release
rknn.release()
cd /workspace/to_docker/project-Toolkit2/
source toolkit2_1.5/bin/activate
cd /workspace/to_docker/rk3588_yolov8/
python pt2rknn.py ./yolov8n_rknnopt.pt rk3588
最后生成的pt2rknn.rknn就可以放到板子使用了。
最后验证能识别自己训练的模型,训练包是在/workspace/to_docker/rk3588_yolov8/datasets/ship1