在windows下使用YOLOv8旋转目标检测算法跑自己的数据集

首先本人参考了以下链接，并使用了相关代码进行记录整合

参考一

1.早茶和猫博客链接：

Yolov8_obb（prob loss) 基于anchor_free的旋转框目标检测，剪枝，跟踪（ByteTracker）_yolov8 obb-CSDN博客

2.早茶和猫代码链接：

GitHub - yzqxy/Yolov8_obb_Prune_Track

参考二

1.Glenn Jocher代码链接：

ultralytics/ultralytics at obb (github.com)

2.ultralytics官方文件代码链接：

Home - Ultralytics YOLOv8 Docs

本人代码链接：

1.根据早茶和猫部分改动后代码：

2.根据Glenn Jocher部分改动后代码：

一.数据集部分

1.官方数据集

(1).DOTA 系列链接：

https://github.com/ultralytics/yolov5/releases/download/v1.0/DOTAv2.zip

2.制作自己的数据集

(1).X-Anylabeling

X-Anylabeling代码链接:CVHub520/X-AnyLabeling: Effortless data labeling with AI support from Segment Anything and other awesome models. (github.com)

X-Anylabeling博客链接:

【2023-2024年最新教程】yolov5_obb: 旋转目标检测从数据制作到终端部署全流程教学_旋转目标检测 数据标注-CSDN博客

(2).Rolabelimg

Rolabelimg代码链接：

cgvict/roLabelImg: Label Rotated Rect On Images for training (github.com)

Rolabelimg博客链接：

【roLabelImg】windows下旋转框标注软件安装、使用、rolabelimg打包成exe_rolabelimg怎么安装-CSDN博客

(3）两种标注格式转换代码

首先两种标注形式本人生成的都是XML文件，感觉在YOLO系列算法上比较方便

X-Anylabeling标注是八点式，即旋转框的四个点的x,y坐标

Rolabeling标注是五点式，即旋转框的中心x,y坐标、旋转框的长度和宽度、旋转角度

五点式XML文件转换为八点式XML文件在转换为YOLO可以训练的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=['你的类别']
def edit_xml(xml_file, dotaxml_file):
    """
    修改xml文件
    :param xml_file:xml文件的路径
    :return:
    """
    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 = r" 已标注并需要转换的xml文件"  
    dotaxml_path = r'存储dota格式的xml文件的输出路径'  #
    out_path = r'存储data格式yolov5_obb可训练的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)

 数据存放格式：

你的数据存放格式
dataset
	--images
		--train
		--val
		--test
	--labelTxt
		--trian
		--val
		--test
	--dataset.yaml #内部存放数据路径、类别名和类别数

将上一步八点式txt数据格式进一步进行归一化转换

此处的数据集转换主要用于Glenn Jocker的代码训练，此处需要下载链接处的代码，然后新建如下所示.py文件

此处如果使用官方数据集，也需要按照此处代码进行转换，如果按照本人给出的链接进行下载，则文件夹中包含已经转换成功的数据集直接进行训练就好

此处如果使用自己的数据集，其中如果是png格式则直接进行训练就好，如果使用的是png以外的其他数据格式，则需要将：

ultralytics.data.coverter中的convert_data_to_yolo_obb处的class_mapping改为自己的数据类别名

同时将以下代码进行注释

# if image_path.suffix != '.png':
#     continue

# 数据转换
from ultralytics.data.converter import convert_dota_to_yolo_obb
convert_dota_to_yolo_obb("dataset")

3.环境准备

两个代码首先都需要按照要求进行：

pip install -r requirements.txt

早茶的猫代码之后需要执行

cd utils/nms_rotated
python setup.py develop

此处会报错，需要将utils\nms_rotated\src\poly_nms_cuda.cu中部分代码进行修改，如下所示

#源码
#define maxn 10
const double eps=1E-8;

__device__ inline int sig(float d){
    return(d>eps)-(d<-eps);
 }

#修改代码
#define maxn 10
// const double eps=1E-8;

__device__ inline int sig(float d){
    return(d>0.00000001)-(d<-0.00000001);

之后再执行

python setup.py develop

Glenn Jocker的代码环境安装直接pip就行不会由太大问题，如果YOLOv8能跑则YOLOv8_obb也可以跑

4.代码训练

按照上述要求修改数据集和环境之后，再新建.py文件进行训练、验证、测试就行，代码如图所示：

早茶和猫系列：

修改tain,val,detect处的数据参数进行训练就行

#训练
python train.py --data DOTAv2/dataset.yaml --hyp data/hyps/obb/hyp.finetune_dota.yaml --cfg models/yaml/yolov8x.yaml   --epochs 5   --batch-size 4   --img 640   --device 0

python train.py --data dataset/dataset.yaml --hyp data/hyps/obb/hyp.finetune_dota.yaml --cfg models/yaml/yolov8x.yaml   --epochs 500   --batch-size 8   --img 640   --device 0

#测试
python detect.py --weights runs/train/exp/weights/best.pt --source dataset/images/val --img 640 --device 0 --conf-thres 0.25 --iou-thres 0.2

#测试--自动标注
#测试图片并保存对应xml文件
python detect_save_xml.py --save-xml --xml_save_path_dir runs/detect/xml --weights runs/train/exp/weights/best.pt --source dataset/images/val/   --img 640 --device 0 --conf-thres 0.25 --iou-thres 0.2

此处代码用于进行测试图像的xml文件保存，保存结果为五点式xml

#测试图片并保存对应xml文件
python detcet_save_xml.py --save-xml --xml_save_path_dir your_xml_save_path/  --weights  runs/train/exp/weights/best.pt   --source dataset/your datafile/images/val/   --img 640 --device 0 --conf-thres 0.25 --iou-thres 0.2 

Glenn Jocker系列：

训练、验证、测试新建文件夹如下：

# 训练
from ultralytics import YOLO
if __name__ == "__main__":
    model = YOLO("yolov8x-obb.yaml")
    model.train(data="dataset/dataset.yaml", epochs=500, imgsz=640, batch=4, workers=4, device=0)


# # 验证
 from ultralytics import YOLO
 if __name__ == "__main__":
   model=Y O LO("J:/MBLL/apple_diameter/rorectangle/yolov8/ultralyticsob-b/runs/obb/train/weights/best.pt")  
   model.val(data="dataset/dataset.yaml")


 # 测试
 from ultralytics import YOLO
 if __name__ == "__main__":
     model = YOLO("J:/MBLL/apple_diameter/rorectangle/yolov8/ultralytics-obb/runs/obb/train/weights/best.pt")
     model.predict(source="dataset/images/val", save=True, save_txt=True)


 # 数据转换
 from ultralytics.data.converter import convert_dota_to_yolo_obb
 convert_dota_to_yolo_obb("dataset")


 # JPG转PNG
 from PIL import Image # 导入Pillow库中的Image模块
 import os # 导入os模块
 # 定义需要转换的文件夹路径
 folder_path = "J:/MBLL/apple_diameter/rorectangle/yolov8/ultralytics-obb/DOTAv1/images/val/"
 # 遍历文件夹中的所有文件
 for filename in os.listdir(folder_path):
     if filename.endswith(".jpg"): # 判断文件是否以.JPG结尾
         # 打开原始图像文件
         img = Image.open(folder_path + filename)
         # 拼接输出文件路径和文件名
         new_filename = os.path.splitext(filename)[0] + ".png"
         new_filepath = os.path.join(folder_path, new_filename)
         # 保存为PNG格式
         img.save(new_filepath, "PNG")
         # 关闭图像文件
         img.close()