YOLO5 旋转模型图片标注及训练

使用工具

图片获取工具—— opencv_display.py
图片标注工具—— rolabelimg
标注转换工具—— roxml_to_dota.py
模型训练代码—— YOLOv5_DOTA_OBB

图片标注

  1. 获取图片
    新建文件夹,用来保存物体图片。图片要求:图片必须为方形(长、宽相等)、图片为 .png 格式。
    下载图片获取工具 opencv_display.py
    修改 opencv_display.py 文件,主要修改起止点坐标以及图片保存目录,具体修改步骤参考视频 获取图片
    执行代码:
    roslaunch kinect2_bridge kinect2_bridge.launch # 打开新终端conda activate YOLOpython opencv_display.py

    如果图片区域内没有要检测的物体,关闭代码重复调整脚本内容,使视野内有物体,在图片显示区域按动键盘 s 键实现图片获取工作。图片获取完成后,按ctrl + c 关闭代码。

  2. 图片标注
    rolabel 虚拟环境下跳转到 ~/Public/roLabelImg 目录下,运行 roLabelImg.py 文件进行图片标注。
    conda activate rolabelcdcd Public/roLabelImg/# 使用前可修改 data/predefined_classes.txt 文件修改要标注的物体类别。python roLabelImg.py

    标注工具界面如下:

    点击界面左侧open dir 按钮,打开保存图片的文件夹。
    点击界面的 Create RotatedRBox 按钮进行旋转标注,拖动鼠标放置标注框,鼠标左击标注框四角的点进行拖动可调整标注框的大小, 鼠标右击标注框四角的点进行拖动可调整标注框的旋转角度,标注完成后点击save 键进行保存,重复以上步骤,完成所有图片的标注,保存时把所有的标注文件保存在同一个文件夹。标注过程参考视频 标注图片

标注处理

  1. 下载训练模型文件。
    下载 YOLOv5_DOTA_OBB 代码到本地,解压文件夹到合适的目录,建议解压到 Documents 目录下。
  2. 下载标注转换工具 roxml_to_dota.py
    在脚本第123行修改 dir 变量,修改为存放标签的文件夹路径,保存,运行该脚本,会在原来的标注文件夹生成 TXT 格式的标注文件。过程参考视频 标注处理_01
  3. 复制图片文件以及标注文件。
    跳转到目录 YOLOv5_DOTA_OBB/DOTA_devkit_YOLO-master 文件夹下,修改文件 dota_utils.py,第18行,修改 classnames_v1_5 内容为物体标注类别,注意类别顺序。
    复制物体图片到 YOLOv5_DOTA_OBB/DOTA_devkit_YOLO-master/DOTA_demo/images 文件夹下,将上个步骤生成的 TXT格式的标注文件复制到 YOLOv5_DOTA_OBB/DOTA_devkit_YOLO-master/DOTA_demo/labelTxt 文件夹下。文件复制完成之后,运行 YOLOv5_DOTA_OBB/DOTA_devkit_YOLO-master/YOLO_Transform.py 文件,运行完成之后会在 YOLOv5_DOTA_OBB/DOTA_devkit_YOLO-master/DOTA_demo/yolo_labels 文件夹下生成新的标注文件。具体步骤参考视频 标注处理_02

模型训练

  1. 复制图片到对应的路径。
    跳转路径到刚才解压的文件夹下,将图片复制到 ./DOTA_demo_view/images 文件夹下,将YOLOv5_DOTA_OBB/DOTA_devkit_YOLO-master/DOTA_demo/yolo_labels 文件夹下新生成的标注文件复制到 ./DOTA_demo_view/labels 文件夹下。
    修改文件 ./data/DOTA_ROTATED.yaml ,修改 nc 为标签类别的数量,name 为标签类别信息,注意标签类别顺序。
  2. 开始训练
    修改 train.py 文件,第652行,将default=[600, 600] 修改为自己图片的分辨率。之后开始训练。如果报显存溢出的错误,可以尝试调整batch-size 以及 device 的值消除报错。
    conda activate YOLO# 安装 swigcd  \.....\yolov5_DOTA_OBB\utilssudo apt-get install swig# 生成 python 库, 该步骤仅执行一次,以后不需要重复执行swig -c++ -python polyiou.ipython setup.py build_ext --inplace# 开始训练cd ..python train.py 
  3. 复制训练好的权重文件。
    训练完成后,终端会输出生成的权重以及权重目录,具体输出如图:

    等待训练完成,将训练好的权重文件(best.pt)复制到目录 /home/ts/catkin_ws/src/yolo_detection/src/YOLO/pick/下。