YOLO5 旋转模型图片标注及训练
使用工具
图片获取工具—— opencv_display.py
图片标注工具—— rolabelimg
标注转换工具—— roxml_to_dota.py
模型训练代码—— YOLOv5_DOTA_OBB
图片标注
- 获取图片
新建文件夹,用来保存物体图片。图片要求:图片必须为方形(长、宽相等)、图片为.png
格式。
下载图片获取工具opencv_display.py
。
修改opencv_display.py
文件,主要修改起止点坐标以及图片保存目录,具体修改步骤参考视频获取图片
。
执行代码:roslaunch kinect2_bridge kinect2_bridge.launch # 打开新终端conda activate YOLOpython opencv_display.py
如果图片区域内没有要检测的物体,关闭代码重复调整脚本内容,使视野内有物体,在图片显示区域按动键盘
s
键实现图片获取工作。图片获取完成后,按ctrl + c
关闭代码。 - 图片标注
在rolabel
虚拟环境下跳转到~/Public/roLabelImg
目录下,运行roLabelImg.py
文件进行图片标注。conda activate rolabelcdcd Public/roLabelImg/# 使用前可修改 data/predefined_classes.txt 文件修改要标注的物体类别。python roLabelImg.py
标注工具界面如下:
点击界面左侧open dir
按钮,打开保存图片的文件夹。
点击界面的Create RotatedRBox
按钮进行旋转标注,拖动鼠标放置标注框,鼠标左击标注框四角的点进行拖动可调整标注框的大小, 鼠标右击标注框四角的点进行拖动可调整标注框的旋转角度,标注完成后点击save
键进行保存,重复以上步骤,完成所有图片的标注,保存时把所有的标注文件保存在同一个文件夹。标注过程参考视频标注图片
。
标注处理
- 下载训练模型文件。
下载YOLOv5_DOTA_OBB
代码到本地,解压文件夹到合适的目录,建议解压到Documents
目录下。 - 下载标注转换工具
roxml_to_dota.py
。
在脚本第123行修改dir
变量,修改为存放标签的文件夹路径,保存,运行该脚本,会在原来的标注文件夹生成 TXT 格式的标注文件。过程参考视频标注处理_01
。 - 复制图片文件以及标注文件。
跳转到目录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
。
模型训练
- 复制图片到对应的路径。
跳转路径到刚才解压的文件夹下,将图片复制到./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
为标签类别信息,注意标签类别顺序。 - 开始训练
修改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
- 复制训练好的权重文件。
训练完成后,终端会输出生成的权重以及权重目录,具体输出如图:
等待训练完成,将训练好的权重文件(best.pt)复制到目录/home/ts/catkin_ws/src/yolo_detection/src/YOLO/pick/
下。