YOLOv11——从零开始制作OBB数据集

冰_gogo

已于 2025-02-17 21:48:53 修改

阅读量2.7k

点赞数 53

文章标签： YOLO python 深度学习计算机视觉

于 2025-02-16 16:30:57 首次发布

本文链接：https://blog.youkuaiyun.com/2301_80344847/article/details/145652347

版权

前提准备：
1、YOLO环境需要配置完成。
2、要准备一张图片，然后利用代码进行角度旋转从而制作出多张图片。代码可私聊获得！从而再开始以下文章的阅读，效果更佳。

一、数据标注

1.1 安装roLabelImg

参考链接: https://blog.youkuaiyun.com/qq_41672428/article/details/107690102 ，完成roLabelImg的安装，熟悉使用。

二、数据集准备

2.1 用roLabelImg打标签，生成原始xml标注文件

在roLabelImg上对所有图片打标签，roLabelImg生成的原始XML标签文件为中心点式，格式为（cx，cy，w，h，angle）。

原始XML标签文件的格式： cx，cy，w，h，angle

标注图片如下：
在这里插入图片描述

生成的原始XML标签文件如下：
在这里插入图片描述

点击上图中第1个xml标签文件，打开后格式如下：
在这里插入图片描述

2.2 将原始的XML标注文件转换为DOTA格式标签文件（TXT）

转换过程：

 第一步：原始XML文件 ——————> DOTA格式的XML文件
  	（ 原始XML文件格式  ：cx，cy，w，h，angle
	—>   DOTA格式的XML文件格式  ：x0，y0，x1，y1，x2，y2，x3，y3 ）
     
 第二步：DOTA格式的XML文件 ——————> DOTA格式的TXT文件
  	（DOTA格式的TXT文件格式 ：x0，y0，x1，y1，x2，y2，x3，y3，cls，cls_index）

【提醒：上述这些标签文件的坐标都是还没有进行归一化的！】

具体实现代码如下：

新建.py文件，然后复制下面代码。（我的名“rolabelImg_XML_to_Dota.py”）
【代码运行后，读取原始XML所在的文件夹里所有标签文件，然后会同时在 DOTA格式的XML和TXT文件夹下生成对应标签文件。】

import os
import xml.etree.ElementTree as ET
import math
"""
代码使用：
    只用修改if __name__ == '__main__': 下面的三条路径即可：
    
    roxml_path = r"path/已经存放着原始XML文件的文件夹名字"  
    dotaxml_path = r"path/准备要存放DOTA格式的XML文件的文件夹名字"  
    			# （小建议：先手动创建该文件夹，然后路径放在这里）
    out_path = r"path/准备要存放DOTA格式的TXT文件的文件夹名字"  
    			# （小建议：先手动创建该文件夹，然后路径放在这里）

"""

def edit_xml(xml_file, dotaxml_file):
    """
    修改 XML 文件，将 bndbox 或 robndbox 转换为四点坐标格式。

    :param  xml_file: 原始 XML 文件路径
    :param  dotaxml_file: 转换后的 XML 文件保存路径
    """
    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")

        if obj.find('robndbox') is None:
            # 处理 bndbox
            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 = float(obj_xmin.text)
            ymin = float(obj_ymin.text)
            xmax = float(obj_xmax.text)
            ymax = float(obj_ymax.text)
            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:
            # 处理 robndbox
            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_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):
    """
    计算旋转后的点坐标。

    :param xc: 旋转中心 x 坐标
    :param yc: 旋转中心 y 坐标
    :param xp: 点 x 坐标
    :param yp: 点 y 坐标
    :param theta: 旋转角度（弧度）
    :return: 旋转后的点坐标 (x, y)
    """
    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 get_unique_classes(xml_path):
    """
    从 XML 文件中提取所有唯一的类别名称。

    :param xml_path: XML 文件所在的目录
    :return: 包含所有唯一类别名称的集合
    """
    unique_classes = set()  # 使用集合存储唯一的类别名称
    files = os.listdir(xml_path)
    for file in files:
        if not file.endswith('.xml'):
            continue
        tree = ET.parse(os.path.join(xml_path, file))
        root = tree.getroot()
        for obj in root.findall('object'):
            cls = obj.find('name').text
            unique_classes.add(cls)
    return unique_classes


def generate_class_mapping(unique_classes):
    """
    根据唯一的类别名称生成类别映射字典。

    :param unique_classes: 包含所有唯一类别名称的集合
    :return: 类别名称到编号的映射字典
    """
    class_mapping = {}
    for index, cls in enumerate(sorted(unique_classes)):  # 按字母顺序排序
        class_mapping[cls] = index
    return class_mapping


def totxt(xml_path, out_path, class_mapping):
    """
    将 XML 文件转换为 TXT 文件，并根据类别名称动态生成编号。

    :param xml_path: XML 文件所在的目录
    :param out_path: 保存 TXT 文件的目录
    :param class_mapping: 类别名称到编号的映射字典
    """
    # 确保输出目录存在
    if not os.path.exists(out_path):
        os.makedirs(out_path)

    # 遍历 xml_path 下的所有文件
    files = os.listdir(xml_path)
    for file in files:
        # 只处理 .xml 文件
        if not file.endswith('.xml'):
            continue

        # 解析 XML 文件
        tree = ET.parse(os.path.join(xml_path, file))
        root = tree.getroot()

        # 获取文件名（不带扩展名）
        name = os.path.splitext(file)[0]

        # 设置输出文件路径
        output = os.path.join(out_path, name + '.txt')  # 使用 os.path.join 确保路径正确
        with open(output, 'w') as f:  # 使用 with 打开文件，确保文件正确关闭
            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))

                # 根据类别名称获取对应的编号
                cls_index = class_mapping.get(cls, -1)  # 如果类别不存在，默认返回 -1
                if cls_index == -1:
                    print(f"Warning: Class '{cls}' not found in class_mapping. Skipping this object.")
                    continue

                # 写入文件
                f.write("{} {} {} {} {} {} {} {} {} {}\n".format(x0, y0, x1, y1, x2, y2, x3, y3, cls, cls_index))
        print(f"Generated: {output}")


if __name__ == '__main__':
    # 设置路径
    roxml_path = r"C:\Users\Users_name\Desktop\pen_label"  # 存放原始 XML 文件夹名字的路径
    dotaxml_path = r"C:\Users\Users_name\Desktop\pen_label_4xml (1)"  # 存放 DOTA格式的 XML文件夹名字（小建议：先手动创建该文件夹，然后路径放在这里）
    out_path = r"C:\Users\Users_name\Desktop\pen_label_4txt (1)"  # 存放 DOTA格式的 TXT 文件夹名字（小建议：先手动创建该文件夹，然后路径放在这里）

    # 第一步：将 XML 文件统一转换成旋转框的 XML 文件
    filelist = os.listdir(roxml_path)
    for file in filelist:
        edit_xml(os.path.join(roxml_path, file), os.path.join(dotaxml_path, file))

    # 第二步：从 XML 文件中提取所有唯一的类别名称
    unique_classes = get_unique_classes(dotaxml_path)
    print(f"Unique classes found: {unique_classes}")

    # 第三步：生成类别映射字典
    class_mapping = generate_class_mapping(unique_classes)
    print(f"Generated class mapping: {class_mapping}")

    # 第四步：将旋转框 XML 文件转换成 TXT 格式
    totxt(dotaxml_path, out_path, class_mapping)

原始XML文件夹以及内容如下：
+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

· 文件夹+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

· 内容

DOTA格式的XML文件夹以及内容如下：
+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

· 文件夹
+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

· 内容

DOTA格式的TXT文件夹以及内容如下：
+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

· 文件夹
+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

· 内容

2.3 划分数据集

【为方便展示，我把下面涉及到的文件夹从桌面搬到
“新建文件夹（2）中进行截图展示。”】
1、这一步，我们需要用到两个文件夹。
第一个文件夹是我们存放所有图片的文件夹（我的名为“pen”）。
第二个文件夹是DOTA格式的TXT标签文件的文件夹名字（我的名为“pen_label_4txt (1)” ）
在这里插入图片描述

2、新建一个空的文件夹（用来当作一个标准数据集格式文件夹）

文件夹格式如下:
 ———Ball Point Pen  （这里只有“Ball Point Pen”可自定义。 images和labels最好保持同名）  
         ├─ images  
         │   ├─ train   （空的，里面的内容由划分数据集的代码所得）
         │   ├─ test    （空的，里面的内容由划分数据集的代码所得）
         │   └─ val     （空的，里面的内容由划分数据集的代码所得）
         └─ labels
             ├─ test    （空的，里面的内容由划分数据集的代码所得）
             ├─ train_original（空的，里面的内容由划分数据集的代码所得）
             └─ val_original  （空的，里面的内容由划分数据集的代码所得）

在这里插入图片描述

3、按照训练集：验证集：测试集=7：2：1的比例进行数据集的划分：
具体实现代码如下

import os
import random
import shutil
from tqdm import tqdm

def split_dataset(
    images_dir: str,  # 图像文件夹路径
    labels_dir: str,  # 标签文件夹路径
    output_root: str,  # 输出根目录
    train_ratio: float = 0.7,  # 训练集比例
    val_ratio: float = 0.2,  # 验证集比例
    image_extensions: tuple = (".jpg", ".png", ".jpeg"),  # 支持的图像文件格式
    seed: int = 42,  # 随机种子
):
    """
    将数据集按照 7:2:1 的比例划分为训练集、验证集和测试集，并将图像和标签文件分别复制到相应的文件夹中。

    参数:
        images_dir (str): 图像文件夹路径。
        labels_dir (str): 标签文件夹路径。
        output_root (str): 输出根目录。
        train_ratio (float): 训练集比例，默认为 0.7。
        val_ratio (float): 验证集比例，默认为 0.2。
        image_extensions (tuple): 支持的图像文件格式，默认为 (".jpg", ".png", ".jpeg")。
        seed (int): 随机种子，默认为 42。
    """
    # 设置随机种子
    random.seed(seed)

    # 创建输出文件夹
    output_images_dir = os.path.join(output_root, "images")
    output_labels_dir = os.path.join(output_root, "labels")

    # 创建训练集、验证集和测试集的图像文件夹
    for subset in ["train", "val", "test"]:
        os.makedirs(os.path.join(output_images_dir, subset), exist_ok=True)  # 图像文件夹

    # 创建训练集和验证集的标签文件夹（指定路径）
    os.makedirs(os.path.join(output_labels_dir, "train_original"), exist_ok=True)  # 训练集标签文件夹
    os.makedirs(os.path.join(output_labels_dir, "val_original"), exist_ok=True)  # 验证集标签文件夹
    os.makedirs(os.path.join(output_labels_dir, "test"), exist_ok=True)  # 测试集标签文件夹

    # 获取所有图像文件的文件名列表
    image_files = [f for f in os.listdir(images_dir) if f.endswith(image_extensions)]
    random.shuffle(image_files)  # 随机打乱文件列表

    # 计算训练集、验证集和测试集的大小
    total_files = len(image_files)
    train_size = int(total_files * train_ratio)  # 训练集大小
    val_size = int(total_files * val_ratio)  # 验证集大小
    test_size = total_files - train_size - val_size  # 测试集大小

    # 复制图像和标签文件到相应的子集文件夹中
    for i, image_file in enumerate(tqdm(image_files, desc="分割数据集中")):
        base_file_name = os.path.splitext(image_file)[0]  # 获取文件名（不包括扩展名）
        image_path = os.path.join(images_dir, image_file)  # 图像文件路径
        label_path = os.path.join(labels_dir, base_file_name + ".txt")  # 标签文件路径

        # 根据索引判断文件应复制到训练集、验证集还是测试集
        if i < train_size:
            # 复制到训练集
            shutil.copy(image_path, os.path.join(output_images_dir, "train", image_file))  # 复制图像
            shutil.copy(label_path, os.path.join(output_labels_dir, "train_original", base_file_name + ".txt"))  # 复制标签
        elif i < train_size + val_size:
            # 复制到验证集
            shutil.copy(image_path, os.path.join(output_images_dir, "val", image_file))  # 复制图像
            shutil.copy(label_path, os.path.join(output_labels_dir, "val_original", base_file_name + ".txt"))  # 复制标签
        else:
            # 复制到测试集
            shutil.copy(image_path, os.path.join(output_images_dir, "test", image_file))  # 复制图像
            shutil.copy(label_path, os.path.join(output_labels_dir, "test", base_file_name + ".txt"))  # 复制标签

    print(f"数据集分割完成！训练集: {train_size} 个样本，验证集: {val_size} 个样本，测试集: {test_size} 个样本。")


# 示例调用
if __name__ == "__main__":
    # 数据集路径
    images_dir = r"C:/Users/Desktop/pen"  # 需要读取的所有图像文件夹路径
    labels_dir = r"C:\Users\Desktop\pen_label_4txt (1)"  # 需要读取的所有图像与之对应的txt标签文件夹路径

    # 输出路径
    output_root = r"C:\Users\Desktop\Ball Point Pen" # 保存最终数据集的根目录

    # 调用函数
    split_dataset(
        images_dir=images_dir,
        labels_dir=labels_dir,
        output_root=output_root,
        train_ratio=0.7,  # 训练集比例
        val_ratio=0.2,  # 验证集比例
        image_extensions=(".jpg", ".png", ".jpeg"),  # 支持的图像文件格式
        seed=42,  # 随机种子
    )

代码运行后，Ball Point Pen文件夹下的所有文件夹里会生成对应的图片和标签。至此，对数据集的划分完成了。

然后，数据集生成后，要复制到官方源代码yolov11用来存放数据集datasets下。
（这里我用来存放数据集datasets的路径：D:\yolo_v8\ultralytics-main\ultralytics-main\datasets。得到 D:\yolo_v8\ultralytics-main\ultralytics-main\datasets\Ball_Point_Pen）

2.4 DOTA格式标签文件（TXT）转为YOLOV11-OBB训练需要数据集格式

虽然数据集实现了划分，但标签文件为DOTA格式的txt文件，不是YOLOV11-OBB训练所需要的txt文件。需要用官方的源代码实现转换。

DOTA格式的txt文件，如下：

+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

YOLOV11-OBB训练所需要的txt文件，如下

怎么利用官方代码实现转换呢？

1、在转换为yolov11-obb训练需要的格式之前，先确保DOTA数据集的文件夹已有如下格式：
（且位置在D:\yolo_v8\ultralytics-main\ultralytics-main\datasets\Ball_Point_Pen中）
在这里插入图片描述

2、然后，我们需要修改ultralytics/data/converter.py中两个地方。
第一步
修改class_map。
将原来的部分注释掉，修改为自己数据集的类别名。
（我这里只有一个类别“Ball_Point_Pen”，因此修改如下：）

第二步
修改image_path.suffix语句。

将源代码中的
if image_path.suffix != “.png”:
修改为：
if image_path.suffix not in [“.png”, “.jpg”, “.jpeg”]:

如下图：
在这里插入图片描述

3、将DOTA格式的TXT标签文件转换为YOLOV11-OBB训练所需的TXT标签文件。

【需要官方源代码根目录下新建.py文件，（我名为“DOTA数据转为yolov8-obb所训练的格式.py”），复制下面代码。即可实现TXT标签文件的转换。代码会在train_original和val_original这两个TXT标签文件的基础上，把转换后的TXT标签文件放在代码自动生成的train和val文件夹中。】

具体实现代码如下：

from ultralytics.data.converter import convert_dota_to_yolo_obb

# 输入需要转换TXT标签文件的数据集文件夹名字的路径
convert_dota_to_yolo_obb(r"..\ultralytics-main\datasets\Ball-Point-Pen")

运行后，数据集的label文件夹下会生成train和val文件夹，如下图：

在这里插入图片描述

提示，在YOLOV11-OBB训练时，label文件夹下的train和val会用到训练中，而train_original和val_original文件不会被用到训练中。

三、配置模型

3.1 配置数据集的yaml文件。

在官方源代码的目录ultralytics-main\ultralytics\cfg\datasets下，新建.yaml文件，名字最好与数据集的名字相同。如下图

（我在D:\yolo_v8\ultralytics-main\ultralytics-main\ultralytics\cfg\datasets，新建一个名为”Ball_Point_Pen.yaml文件)
在这里插入图片描述

3.2 配置yolov11-obb.yaml文件。

修改官方源代码中ultralytics-main\ultralytics\cfg\models\11\yolo11-obb.yaml中类别数量nc，数量修改为自己OBB数据集的类别数量。如下图：

（我这里是D:\yolo_v8\ultralytics-main\ultralytics-main\ultralytics\cfg\models\11\yolo11-obb.yaml，只有1个类别“Ball_Point_Pen”）
在这里插入图片描述

四、模型训练

4.1 训练

在存放OBB数据集的路径下，新建.py文件用于训练。

（这里我在D:\yolo_v8\ultralytics-main\ultralytics-main\datasets\Ball_Point_Pen下，新建train.py文件）
在这里插入图片描述

具体代码如下:

from ultralytics import YOLO


def main():
    model = YOLO(r'D:\ultralytics-main\ultralytics\cfg\models\11\yolo11n-obb.yaml').load(r'D:\ultralytics-main\yolo11n-obb.pt')  # build from YAML and transfer weights
    model.train(data=r'D:\ultralytics-main\ultralytics\cfg\datasets\Ball-Point-Pen.yaml',
                epochs=200,
                imgsz=640,
                batch=16,
                workers=2)


if __name__ == '__main__':
    main()

五、模型预测

5.1 预测

在存放OBB数据集的路径下，新建.py文件用于预测。

（这里我在D:\yolo_v8\ultralytics-main\ultralytics-main\datasets\Ball-Point-Pen下，新建predict

.py文件）
在这里插入图片描述

具体代码如下:

from ultralytics import YOLO

# model = YOLO('path/to/best.pt')  # load a custom model
model = YOLO(r'D:\ultralytics-main\datasets\Ball-Point-Pen\runs\obb\train\weights\best.pt')  # load an official model

# Predict with the model
# results = model('https://ultralytics.com/images/bus.jpg')  # predict on an image
results = model(r"D:\ultralytics-main\datasets\Ball-Point-Pen\images\test\065（65度）.jpg",
                save=True)