【记录】使用yolov5_obb训练自己的数据集

引言

对于寻常的yolov5目标检测任务，只能检测水平或者垂直的检测框，而对于旋转框的检测却无能为力。为此，在这记录下使用yolov5_obb来训练自己数据集。

一、准备数据集

1、我们先看所需要的数据集文件什么样子，如下图文件夹Symbol所示：

images文件夹下是格式为PNG的训练图片，labelTxt下则是对应的txt文件。也就是说我们只需要获取txt文件就行。 

2、标注图像

为了能够标注旋转的目标检测框，我们采用rolabelImg工具进行标注，具体教程在我以前的文章有（点个赞收藏下呗）：

【教程】安装labelme以及旋转标记工具roLabelimg-CSDN博客文章浏览阅读69次。因为需要标记旋转的目标检测框，所以要用到工具roLabelimg，labelme则更多用来标注图像分割，并且使用roLabelimg的前提是安装labelme。（window环境下）https://blog.csdn.net/qq_39149619/article/details/133987812?spm=1001.2014.3001.55013、xml转换txt

步骤2得到的图像是xml格式的，而训练则需要txt文件，博主在这里给出转换代码如下：

import os
import math
import xml.etree.ElementTree as ET
import numpy as np
import cv2

def rotate(cx, cy, w, h, angle):
    # 用于常规坐标时，angle是顺时针旋转角度
    # 用于图像坐标时，angle是逆时针旋转角度
    # math.cos(angle)这里的angle单位是rad
    # angle = angle/180*math.pi  单位转换
    angle=-angle#这里是图像坐标
    points = [[cx-w/2, cy-h/2], [cx+w/2, cy-h/2],
              [cx+w/2, cy+h/2], [cx-w/2, cy+h/2]]
    newpoints = []
    if angle < 0:  # 逆时针
        angle = -angle
        for point in points:
            x, y = point
            newx = round((x-cx)*math.cos(angle) - (y-cy)*math.sin(angle) + cx,1)
            newy = round((x-cx)*math.sin(angle) + (y-cy)*math.cos(angle) + cy,1)
            newpoints.append([newx, newy])
    else:
        for point in points:
            x, y = point
            newx = round((x-cx)*math.cos(angle) + (y-cy)*math.sin(angle) + cx,1)
            newy = round((y-cy)*math.cos(angle) - (x-cx)*math.sin(angle) + cy,1)
            newpoints.append([newx, newy])
    return newpoints

def roxml2txt(dir):
    #dir是xml文件目录
    files = os.listdir(dir)
    parentdir,dirname = os.path.split(dir)
    txtdir=os.path.join(parentdir,'txt')
    if not os.path.exists(txtdir):
        os.mkdir(txtdir)
    for f in files:
        if f == 'desktop.ini':
            continue
        xml_file = open(os.path.join(dir,f), encoding='gbk')
        root = ET.parse(xml_file).getroot()
        # xml = ET.parse(os.path.join(dir,f))
        # root = xml.getroot()
        boxes = root.iter('robndbox')
        with open(os.path.join(txtdir,f.split('.')[0]+'.txt'),'w+') as t:
            for box in boxes:
                cx = float(box.find('cx').text)
                cy = float(box.find('cy').text)
                w = float(box.find('w').text)
                h = float(box.find('h').text)
                angle = float(box.find('angle').text)
                newpoints = rotate(cx, cy, w, h, angle)#计算旋转后的4个点坐标
                # 用于查看坐标转换是否正确，在原图上画矩形框，自行修改图片路径
                newpoints=np.array(newpoints)
                newpoints= newpoints.astype(int)
                img=cv2.imread(os.path.join('test','images',f.split('.')[0]+'.png'))
                img=cv2.polylines(img,[newpoints],isClosed=True,color=(0,0,255))
                cv2.imshow('pic',img)
                cv2.waitKey()
                ##########################################
                line=''
                for point in newpoints:
                    line+=str(point[0])+' '+str(point[1])+' '
                line+='word 0\n'
                t.write(line)
                print(line)
            t.close()

roxml2txt(r'F:\pycharm\person_AI_astional\Script\label_xml')

注意修改自己xml文件路径，以及不要忘记建test文件（仔细看代码）。

二、准备代码

1、经过步骤一，我们已经得到能够训练的数据集了（PNG和txt），代码参考的是别的大佬的，这里给出链接：hukaixuan19970627/yolov5_obb: yolov5 + csl_label.(Oriented Object Detection)（Rotation Detection）（Rotated BBox）基于yolov5的旋转目标检测 (github.com)

下载好解压至任意目录，安装所需环境即可。

2、 将数据集放入dataset文件夹下，设置数据集的读取路径。打开data/yolov5obb_demo.yaml文件，修改数据集读取路径（为了方便这里验证集即是训练集，具体根据自己需求更改）

# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
path: ./dataset # dataset root dir
train: Symbol/images #images   # train images (relative to 'path')
val: Symbol/images #images  # val images (relative to 'path')
test: Symbol/images  #images # test images (optional)

# Classes
nc: 1  # number of classes
names: ['symbol']  # class names

3、其他没什么需要修改的，打开train.py文件，执行以下的命令开始训练(具体的参数还得根据你自己需求来)：

python train.py --weights weights/yolov5s.pt --data data/yolov5obb_demo.yaml --epochs 100 --batch-size 16 --imgsz 840 --adam

三、遇到的问题以及解决方法

1、花时间最多的就是在安装yolov5_obb环境上了，我这里给出一点经验：建议完全遵守作者所给install文件中的安装方法，仔细检查CUDA Driver Version和 CUDA Toolkit Version的版本，并且 CUDA Toolkit Version要小于等于CUDA Driver Version！

说到头也就是你nvidia-smi出来的版本要大于等于nvcc -V出来的版本。另外，python版本建议3.9，不然按照作者那个安装会导致torch下载不下来。

很多人会在python setup.py develop安装nms_rotated出错，基本上都是CUDA Driver Version和 CUDA Toolkit Version两者版本出现问题，仔细检查。另外，pip下载的cu11.3版本，确保CUDA Toolkit Version大于等于11.3，以下是博主的版本：

2、当我们能够正确安装好环境之后，信心满满的执行train.py。结果如下图，chao：

错误说什么labels, shapes, self.segments = zip(*cache.values()) ValueError: not enough values to unpack (expected 3, got 0)

解决方法：是Numpy版本太高的原因，我们执行下面语句降低版本：

pip install numpy==1.22.2

最后放一下博主的环境版本信息：

还有其他问题可以评论里留言，我都会解答的！

别忘记收藏点赞！