RK3588 YOLOV8训练转换流程

训练

数据集最开始是下面的结构

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