制作自己的旋转框数据集来训练YOLOv11并使用

啊某十四

已于 2024-11-06 10:54:28 修改

阅读量2.5k

点赞数 48

文章标签： YOLO python

于 2024-11-05 19:05:26 首次发布

本文链接：https://blog.csdn.net/qq_56706664/article/details/143501783

版权

一、旋转框标注

制作旋转框的标注我这里用到的是rolabelimg，因为之前已经下载了，没下载过的同学可以去自行下载，打开anaconda prompt（anaconda也得下载，使用它管理环境很方便）

激活我们创建的rolabelimg的环境：

conda activate rolabelimg #这里我的rolabelimg是我的环境名

然后跳转到文件所在的本地地址

直接python运行该文件：

python roLabelImg.py

然后就会出现标注界面啦

首先，我们点击选好我们要标注的文件夹；

然后可以选择两种检测框进行标注，上面的是标准框，下面的是旋转框，想要什么类型的检测框点击然后标注就行，因为现在做obb检测，所以我们选择旋转框来进行标注，在标注目标的大概区域画好旋转框，会弹出提示框要求标注该框的标签名

接着进行微调，鼠标选中旋转框的一个角，右键是旋转角度，左键是缩放，标注完是这样的

注意，每标好一张图片就得点击保存一下标注文件，我一般是在图片对应的文件夹里创建一个labels存放这些标签文件

保存生成的标签文件长这样

二、数据转化

接下来就是把标注好的数据转换成YOLO-obb可以识别的标注数据了。因为我借鉴的是遥感卫星的旋转框技术，所以先把原文件转化成DOTA旋转框的xml文件，再转化成DOTA的txt文件，转化代码如下：

# 文件名称   ：roxml_to_dota.py
# 功能描述   ：把rolabelimg标注的xml文件转换成dota能识别的xml文件，
#             再转换成dota格式的txt文件
#            把旋转框 cx,cy,w,h,angle，或者矩形框cx,cy,w,h,转换成四点坐标x1,y1,x2,y2,x3,y3,x4,y4
import os
import xml.etree.ElementTree as ET
import math

cls_list = ['dlp']  # 一定要改成自己标注时写的标签


def edit_xml(xml_file, dotaxml_file):
    """
    修改xml文件
    :param xml_file:xml文件的路径
    :return:
    """

    # dxml_file = open(xml_file,encoding='gbk')
    # tree = ET.parse(dxml_file).getroot()

    tree = ET.parse(xml_file)
    objs = tree.findall('object')
    for ix, obj in enumerate(objs):
        x0 = ET.Element("x0")  # 创建节点
        y0 = ET.Element("y0")
        x1 = ET.Element("x1")
        y1 = ET.Element("y1")
        x2 = ET.Element("x2")
        y2 = ET.Element("y2")
        x3 = ET.Element("x3")
        y3 = ET.Element("y3")
        # obj_type = obj.find('bndbox')
        # type = obj_type.text
        # print(xml_file)

        if (obj.find('robndbox') == None):
            obj_bnd = obj.find('bndbox')
            obj_xmin = obj_bnd.find('xmin')
            obj_ymin = obj_bnd.find('ymin')
            obj_xmax = obj_bnd.find('xmax')
            obj_ymax = obj_bnd.find('ymax')
            # 以防有负值坐标
            xmin = max(float(obj_xmin.text), 0)
            ymin = max(float(obj_ymin.text), 0)
            xmax = max(float(obj_xmax.text), 0)
            ymax = max(float(obj_ymax.text), 0)
            obj_bnd.remove(obj_xmin)  # 删除节点
            obj_bnd.remove(obj_ymin)
            obj_bnd.remove(obj_xmax)
            obj_bnd.remove(obj_ymax)
            x0.text = str(xmin)
            y0.text = str(ymax)
            x1.text = str(xmax)
            y1.text = str(ymax)
            x2.text = str(xmax)
            y2.text = str(ymin)
            x3.text = str(xmin)
            y3.text = str(ymin)
        else:
            obj_bnd = obj.find('robndbox')
            obj_bnd.tag = 'bndbox'  # 修改节点名
            obj_cx = obj_bnd.find('cx')
            obj_cy = obj_bnd.find('cy')
            obj_w = obj_bnd.find('w')
            obj_h = obj_bnd.find('h')
            obj_angle = obj_bnd.find('angle')
            cx = float(obj_cx.text)
            cy = float(obj_cy.text)
            w = float(obj_w.text)
            h = float(obj_h.text)
            angle = float(obj_angle.text)
            obj_bnd.remove(obj_cx)  # 删除节点
            obj_bnd.remove(obj_cy)
            obj_bnd.remove(obj_w)
            obj_bnd.remove(obj_h)
            obj_bnd.remove(obj_angle)

            x0.text, y0.text = rotatePoint(cx, cy, cx - w / 2, cy - h / 2, -angle)
            x1.text, y1.text = rotatePoint(cx, cy, cx + w / 2, cy - h / 2, -angle)
            x2.text, y2.text = rotatePoint(cx, cy, cx + w / 2, cy + h / 2, -angle)
            x3.text, y3.text = rotatePoint(cx, cy, cx - w / 2, cy + h / 2, -angle)

        # obj.remove(obj_type)  # 删除节点
        obj_bnd.append(x0)  # 新增节点
        obj_bnd.append(y0)
        obj_bnd.append(x1)
        obj_bnd.append(y1)
        obj_bnd.append(x2)
        obj_bnd.append(y2)
        obj_bnd.append(x3)
        obj_bnd.append(y3)

        tree.write(dotaxml_file, method='xml', encoding='utf-8')  # 更新xml文件


# 转换成四点坐标
def rotatePoint(xc, yc, xp, yp, theta):
    xoff = xp - xc;
    yoff = yp - yc;
    cosTheta = math.cos(theta)
    sinTheta = math.sin(theta)
    pResx = cosTheta * xoff + sinTheta * yoff
    pResy = - sinTheta * xoff + cosTheta * yoff
    return str(int(xc + pResx)), str(int(yc + pResy))


def totxt(xml_path, out_path):
    # 想要生成的txt文件保存的路径，这里可以自己修改

    files = os.listdir(xml_path)
    i = 0
    for file in files:

        tree = ET.parse(xml_path + os.sep + file)
        root = tree.getroot()

        name = file.split('.')[0]

        output = out_path + '\\' + name + '.txt'
        file = open(output, 'w')
        i = i + 1
        objs = tree.findall('object')
        for obj in objs:
            cls = obj.find('name').text
            box = obj.find('bndbox')
            x0 = int(float(box.find('x0').text))
            y0 = int(float(box.find('y0').text))
            x1 = int(float(box.find('x1').text))
            y1 = int(float(box.find('y1').text))
            x2 = int(float(box.find('x2').text))
            y2 = int(float(box.find('y2').text))
            x3 = int(float(box.find('x3').text))
            y3 = int(float(box.find('y3').text))
            if x0 < 0:
                x0 = 0
            if x1 < 0:
                x1 = 0
            if x2 < 0:
                x2 = 0
            if x3 < 0:
                x3 = 0
            if y0 < 0:
                y0 = 0
            if y1 < 0:
                y1 = 0
            if y2 < 0:
                y2 = 0
            if y3 < 0:
                y3 = 0
            for cls_index, cls_name in enumerate(cls_list):
                if cls == cls_name:
                    file.write("{} {} {} {} {} {} {} {} {} {}\n".format(x0, y0, x1, y1, x2, y2, x3, y3, cls, cls_index))
        file.close()
        # print(output)
        print(i)


if __name__ == '__main__':
    # -----**** 第一步：把xml文件统一转换成旋转框的xml文件 ****-----
    roxml_path = 'rolabelimg得到的标注文件夹所在路径'
    dotaxml_path = '你存放dota的xml文件所期望的路径'
    out_path = '你存放dota的txt文件所期望的路径'
    filelist = os.listdir(roxml_path)
    for file in filelist:
        edit_xml(os.path.join(roxml_path, file), os.path.join(dotaxml_path, file))

    # -----**** 第二步：把旋转框xml文件转换成txt格式 ****-----
    totxt(dotaxml_path, out_path)

可以创建一个data_transfor文件夹这样好管理，里面再创建三个文件夹分别存放org_xml、dota_xml、dota_txt，然后把标注好的标签数据放进org_xml，然后运行代码，打开dota_txt会得到

打开里面的文件查看是否有数据，如果显示这样

那就转化成功了，就可以下一步了。

接着我们得转化成yolo自己的txt格式，可以用到yolo自带的函数convert_dota_to_yolo_obb，YOLO代码下载跟使用，应该有很多大佬早就写了很多详细的了，这里就不过多赘述了

首先我们得搭建符合该函数的datasets，

接着我们运行代码执行convert_dota_to_yolo_obb函数

from ultralytics.data.converter import convert_dota_to_yolo_obb
convert_dota_to_yolo_obb(r'D:\PycharmProjects\ultralytics-obb\dataobb')

报错：

不要慌，这里就是显示我们的标签名跟他函数里的标签名不一样，我们进到该函数里

然后在运行那个函数，没有报错，转化好的yolo-obb格式的文件就在labels的train跟val里面啦，数据就转化好了。格式是这样的

三、YOLO-obb训练

数据准备好了接下来就是训练了，首先，我们去到官网下载obb的预训练权重

这里我专门建了一个weights文件夹专门放权重文件，我们就把下载好的权重放进weights文件夹里。

yolo11-obb.yaml里的nc改成自己要识别的类数量，我这里只做了单类的识别，所以改成1；

我的数据集的配置文件是这样的，按照配置文件自己调好数据集的位置关系，不一定非得这个搭配，您怎么方便管理怎么来，这里的配置文件只需要你在这个路径下有images跟对应的labels就可以，然后设置类名，我这里只有一个类dlp所以就是0： dlp.

所有步骤都添加好了，我们就可以开始训练了，训练代码如下:

from ultralytics import YOLO


def main():
    model = YOLO("ultralytics/cfg/models/11/yolo11-obb.yaml").load("weights/yolo11n-obb.pt")
    model.train(data="ultralytics/cfg/datasets/coco128.yaml", epochs=200, batch=4, workers=1)


if __name__ == '__main__':
    main()

下图这样就显示训练好了，训练结果在runs\obb\train3中