yolo11n-obb训练rknn模型
必备:
准备一台ubuntu22的服务器或者虚拟机(x86_64)
1、数据集标注:
1)推荐使用X-AnyLabeling标注工具
2)标注选【旋转框】
3)可选AI标注,再手动补充,提高标注速度
4)导出->导出yolo旋转框标签->选择一个class.txt文件(里面写你标注的标签名)
2、下载环境
使用rknn修改后的ultralytics_yolo11项目:ultralytics_yolo11
ONNX转换为RKNN模型需要使用官方rknn_model_zoo工具:rknn_model_zoo-2.2.0
官方rknn-toolkit2工具:rknn-toolkit2-2.2.0
*** 使用git clone项目后,可以使用下面命令切换到对应分支
git checkout v2.2.0
3、安装环境
1)进入ultralytics_yolo11目录,安装依赖。
把下面内容保存到requirement.txt
# Ultralytics requirements
# Usage: pip install -r requirements.txt# Base ----------------------------------------
matplotlib>=3.2.2
numpy>=1.18.5
opencv-python>=4.6.0
Pillow>=7.1.2
PyYAML>=5.3.1
requests>=2.23.0
scipy>=1.4.1
torch>=1.7.0
torchvision>=0.8.1
tqdm>=4.64.0# Logging -------------------------------------
tensorboard>=2.4.1
# clearml
# comet# Plotting ------------------------------------
pandas>=1.1.4
seaborn>=0.11.0# Export --------------------------------------
# coremltools>=6.0 # CoreML export
# onnx>=1.12.0 # ONNX export
# onnx-simplifier>=0.4.1 # ONNX simplifier
# nvidia-pyindex # TensorRT export
# nvidia-tensorrt # TensorRT export
# scikit-learn==0.19.2 # CoreML quantization
# tensorflow>=2.4.1 # TF exports (-cpu, -aarch64, -macos)
# tensorflowjs>=3.9.0 # TF.js export
# openvino-dev>=2022.3 # OpenVINO export# Extras --------------------------------------
ipython # interactive notebook
psutil # system utilization
thop>=0.1.1 # FLOPs computation
# albumentations>=1.0.3
# pycocotools>=2.0.6 # COCO mAP
# roboflow
然后安装
pip install -r requirement.txt
2)安装rknn-toolkit2
注意:310对应的是Python3.10版本 根据自己的python版本选择。支持 python 3.6 - - 3.12版本
cd ~/rknn-toolkit2/rknn-toolkit2/packages
pip install -r requirements_cp310-2.2.0.txt
pip install rknn_toolkit2-2.2.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
4、训练模型
1)保存为dataset.yaml
train: /home/admin/labels
val: /home/admin/labels
nc: 1
names: ['cat']
2)修改ultralytics/cfg/models/11/yolo11-obb.yaml(注意 nc值要和上面names数组的数量一致)
nc: 1 # number of classes
3)保存为train.py,执行python train.py开始训练
from ultralytics import YOLO
import torchmodel_yaml_path = "ultralytics/cfg/models/11/yolo11-obb.yaml"
#数据集配置文件
data_yaml_path = r'dataset.yaml'def main():#torch.backends.cudnn.enabled = Falsemodel = YOLO(model=model_yaml_path) # build from YAML and transfer weightsmodel.info()model.train(data=data_yaml_path,epochs=100,imgsz=960,batch=10,amp=False,workers=2,degrees=180.0)if __name__ == '__main__':main()
5)训练完成后,获得模型地址
/home/admin/ultralytics_yolo11/runs/obb/train/weights/best.pt
5、导出rknn模型
1)修改ultralytics/cfg/default.yaml
diff --git a/ultralytics/cfg/default.yaml b/ultralytics/cfg/default.yaml
index 97f7239e..ca648291 100644
--- a/ultralytics/cfg/default.yaml
+++ b/ultralytics/cfg/default.yaml
@@ -5,13 +5,13 @@ task: detect # (str) YOLO task, i.e. detect, segment, classify, pose, obbmode: train # (str) YOLO mode, i.e. train, val, predict, export, track, benchmark# Train settings -------------------------------------------------------------------------------------------------------
-model: yolo11n.pt # (str, optional) path to model file, i.e. yolo11n.pt, yolo11n.yaml
+model: /home/admin/ultralytics_yolo11/runs/obb/train/weights/best.pt # (str, optional) path to model file, i.e. yolo11n.pt, yolo11n.yamldata: # (str, optional) path to data file, i.e. coco8.yamlepochs: 100 # (int) number of epochs to train fortime: # (float, optional) number of hours to train for, overrides epochs if suppliedpatience: 100 # (int) epochs to wait for no observable improvement for early stopping of trainingbatch: 16 # (int) number of images per batch (-1 for AutoBatch)
-imgsz: 640 # (int | list) input images size as int for train and val modes, or list[h,w] for predict and export modes
+imgsz: 960 # (int | list) input images size as int for train and val modes, or list[h,w] for predict and export modessave: True # (bool) save train checkpoints and predict resultssave_period: -1 # (int) Save checkpoint every x epochs (disabled if < 1)cache: False # (bool) True/ram, disk or False. Use cache for data loading
执行
export PYTHONPATH=./
python ./ultralytics/engine/exporter.py
得到
/home/admin/ultralytics_yolo11/runs/obb/train/weights/best.onnx
2)onnx转换为rknn
修改 ~/rknn_model_zoo/examples/yolo11/python/yolo11.py
diff --git a/examples/yolo11/python/yolo11.py b/examples/yolo11/python/yolo11.py
index 0f8f19c..d77e98d 100644
--- a/examples/yolo11/python/yolo11.py
+++ b/examples/yolo11/python/yolo11.py
@@ -19,7 +19,7 @@ NMS_THRESH = 0.45# The follew two param is for map test# OBJ_THRESH = 0.001# NMS_THRESH = 0.65
-
+"""IMG_SIZE = (640, 640) # (width, height), such as (1280, 736)CLASSES = ("person", "bicycle", "car","motorbike ","aeroplane ","bus ","train","truck ","boat","traffic light",
@@ -33,7 +33,11 @@ CLASSES = ("person", "bicycle", "car","motorbike ","aeroplane ","bus ","train","coco_id_list = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 27, 28, 31, 32, 33, 34,35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63,64, 65, 67, 70, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 84, 85, 86, 87, 88, 89, 90]
+"""+IMG_SIZE = (960, 960)
+CLASSES = ("cat")
+coco_id_list = [0]def filter_boxes(boxes, box_confidences, box_class_probs):"""Filter boxes with object threshold.
开始转换
python convert.py /home/admin/ultralytics_yolo11/runs/obb/train/weights/best.onnx rk3566
得到
/home/admin/rknn_model_zoo/examples/yolo11/model/yolo11.rknn
yolo11.rknn就是yolo11-obb.rknn模型。
用~/rknn_model_zoo/examples/yolov8_obb项目测试,可以正常使用