YOLOv5_OBB 入门指南：从数据准备到模型训练与推理

金畏战Goddard

于 2025-06-27 09:03:34 发布

阅读量272

点赞数 4

CC 4.0 BY-SA版权

本文链接：https://blog.youkuaiyun.com/gitblog_00038/article/details/148941143

YOLOv5_OBB 入门指南：从数据准备到模型训练与推理

yolov5_obb yolov5 + csl_label.(Oriented Object Detection)（Rotation Detection）（Rotated BBox）基于yolov5的旋转目标检测项目地址: https://gitcode.com/gh_mirrors/yo/yolov5_obb

项目概述

YOLOv5_OBB 是基于 YOLOv5 框架开发的面向旋转目标检测（Oriented Object Detection）的开源项目。与标准 YOLOv5 不同，该项目专门针对需要检测带角度信息的物体场景，如遥感图像中的飞机、车辆等目标检测。本文将详细介绍如何使用 YOLOv5_OBB 进行模型训练和推理。

数据准备

1. 标注格式要求

YOLOv5_OBB 要求使用多边形标注格式，具体格式如下：

x1 y1 x2 y2 x3 y3 x4 y4 classname difficult

其中：

(x1,y1) 到 (x4,y4) 表示目标的四个顶点坐标
classname 是目标类别名称
difficult 表示目标检测难度（通常设为0）

示例标注：

1686.0 1517.0 1695.0 1511.0 1711.0 1535.0 1700.0 1541.0 large-vehicle 1

2. 数据集组织结构

推荐的数据集目录结构如下：

父目录
├── yolov5_obb
└── datasets
    └── DOTAv1.5
        ├── train_split_rate1.0_subsize1024_gap200
        ├── val_split_rate1.0_subsize1024_gap200
        └── test_split_rate1.0_subsize1024_gap200
            ├── images
                 |────1.jpg
                 |────...
                 └────10000.jpg
            ├── labelTxt
                 |────1.txt
                 |────...
                 └────10000.txt

3. 高分辨率图像处理

对于高分辨率遥感图像（如DOTA数据集），建议先进行图像分割处理：

python DOTA_devkit/ImgSplit_multi_process.py

这一步骤会将大尺寸图像分割为多个小尺寸子图，有助于提升模型训练效果和推理速度。

模型训练

1. 单GPU训练

使用指定GPU进行训练（示例使用GPU 3）：

python train.py --device 3

2. 多GPU分布式训练

使用4个GPU进行分布式数据并行训练：

python -m torch.distributed.launch --nproc_per_node 4 train.py --device 0,1,2,3

3. 完整训练示例

3.1 原始数据集训练

python train.py \
  --weights 'weights/yolov5n_s_m_l_x.pt' \
  --data 'data/yolov5obb_demo.yaml' \
  --hyp 'data/hyps/obb/hyp.finetune_dota.yaml' \
  --epochs 10 \
  --batch-size 1 \
  --img 1024 \
  --device 0

3.2 分割后数据集训练

python train.py \
  --weights 'weights/yolov5n_s_m_l_x.pt' \
  --data 'data/yolov5obb_demo_split.yaml' \
  --hyp 'data/hyps/obb/hyp.finetune_dota.yaml' \
  --epochs 10 \
  --batch-size 2 \
  --img 1024 \
  --device 0

关键参数说明：

--weights: 预训练模型权重路径
--data: 数据集配置文件路径
--hyp: 超参数配置文件路径
--epochs: 训练轮数
--batch-size: 批次大小
--img: 输入图像尺寸

模型推理与评估

1. 分割数据集推理流程

1.1 获取水平框(HBB)指标

python val.py \
 --data 'data/yolov5obb_demo_split.yaml' \
 --weights 'runs/train/yolov5m_csl_dotav1.5/weights/best.pt' \
 --batch-size 2 --img 1024 --task 'val' --device 0 --save-json --name 'obb_demo_split'

1.2 转换为多边形格式结果

python tools/TestJson2VocClassTxt.py \
  --json_path 'runs/val/obb_demo_split/best_obb_predictions.json' \
  --save_path 'runs/val/obb_demo_split/obb_predictions_Txt'

1.3 合并分割结果

python DOTA_devkit/ResultMerge_multi_process.py \
  --scrpath 'runs/val/obb_demo_split/obb_predictions_Txt' \
  --dstpath 'runs/val/obb_demo_split/obb_predictions_Txt_Merged'

1.4 获取旋转框(OBB)指标

python DOTA_devkit/dota_evaluation_task1.py \
  --detpath 'runs/val/obb_demo_split/obb_predictions_Txt_Merged/Task1_{:s}.txt' \
  --annopath 'dataset/dataset_demo/labelTxt/{:s}.txt' \
  --imagesetfile 'dataset/dataset_demo/imgnamefile.txt'

2. 原始数据集推理流程

2.1 获取水平框(HBB)指标

python val.py \
 --data 'data/yolov5obb_demo.yaml' \
 --weights 'runs/train/yolov5m_csl_dotav1.5/weights/best.pt' \
 --batch-size 1 --img 2048 --task 'val' --device 0 --save-json --name 'obb_demo'

2.2 转换为多边形格式结果

python tools/TestJson2VocClassTxt.py \
  --json_path 'runs/val/obb_demo/best_obb_predictions.json' \
  --save_path 'runs/val/obb_demo/obb_predictions_Txt'

2.3 获取旋转框(OBB)指标

python DOTA_devkit/dota_evaluation_task1.py \
  --detpath 'runs/val/obb_demo/obb_predictions_Txt/Task1_{:s}.txt' \
  --annopath 'dataset/dataset_demo/labelTxt/{:s}.txt' \
  --imagesetfile 'dataset/dataset_demo/imgnamefile.txt'

实际应用推理

对图像、视频等进行推理检测：

python detect.py \
  --weights 'runs/train/yolov5m_csl_dotav1.5/weights/best.pt' \
  --source 'dataset/dataset_demo/images/' \
  --img 2048 --device 0 --conf-thres 0.25 --iou-thres 0.2 \
  --hide-labels --hide-conf

参数说明：