yolov8训练旋转目标检测obb数据集

目录

一、代码

二、数据集转换(VOC 转 data 转 yolo)

1、标注生成原始数据集

2、voc转data

3、划分数据集

4、将labels的data格式转换成yolo格式

三、训练


一、代码

直接用官方的代码

官方源码

二、数据集转换(VOC 转 data 转 yolo)

1、标注生成原始数据集

经过rolabelimg标注后的数据集格式。可以直接将Annotations和JPEGImages两个文件夹放在新建的datasets下面。也可以和我一样将detect和obb做个分类。

2、voc转data

在项目根目录下创建一个roxml_to_dota.py文件,用于转换成dota能用的obb的txt文件。

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


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 = 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:
            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)
    for file in files:

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

        name = file.strip('.xml')
        output = out_path + name + '.txt'
        file = open(output, 'w')

        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))
            file.write("{} {} {} {} {} {} {} {} {} 0\n".format(x0, y0, x1, y1, x2, y2, x3, y3, cls))
        file.close()
        print(output)


if __name__ == '__main__':
    # -----**** 第一步:把xml文件统一转换成旋转框的xml文件 ****-----
    roxml_path = "./datasets/obb/Annotations/"  # 目录下保存的是需要转换的xml文件
    dotaxml_path = './datasets/obb/dotaxml/'
    out_path = './datasets/obb/labelTxt/'
    filelist = os.listdir(roxml_path)
    for path in (dotaxml_path, out_path):
        if not os.path.exists(path):
            os.makedirs(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)

运行完后生成两个文件夹

3、划分数据集

在项目根目录下新建一个splitdata.py文件,

# 将图片和标注数据按比例切分为 训练集和测试集
import shutil
import random
import os
import argparse
 
 
# 检查文件夹是否存在
def mkdir(path):
    if not os.path.exists(path):
        os.makedirs(path)
 
def main(image_dir, txt_dir, save_dir):
    # 创建文件夹
    mkdir(save_dir)
    images_dir = os.path.join(save_dir, 'images')
    labels_dir = os.path.join(save_dir, 'labels')
 
    img_train_path = os.path.join(images_dir, 'train')
    img_val_path = os.path.join(images_dir, 'val')
 
    label_train_path = os.path.join(labels_dir, 'train_original')
    label_val_path = os.path.join(labels_dir, 'val_original')
 
    mkdir(images_dir)
    mkdir(labels_dir)
    mkdir(img_train_path)
    mkdir(img_val_path)
    mkdir(label_train_path)
    mkdir(label_val_path)
 
    # 数据集划分比例,训练集75%,验证集15%,测试集15%,按需修改
    train_percent = 0.50
    val_percent = 0.50
 
    total_txt = os.listdir(txt_dir)
    num_txt = len(total_txt)
    list_all_txt = range(num_txt)  # 范围 range(0, num)
 
    num_train = int(num_txt * train_percent)
    num_val = int(num_txt * val_percent)
 
    train = random.sample(list_all_txt, num_train)
    # 在全部数据集中取出train
    val = [i for i in list_all_txt if not i in train]
    # 再从val_test取出num_val个元素,val_test剩下的元素就是test
    # val = random.sample(list_all_txt, num_val)
 
    print("训练集数目:{}, 验证集数目:{}".format(len(train), len(val)))
    for i in list_all_txt:
        name = total_txt[i][:-4]
 
        srcImage = os.path.join(image_dir, name + '.jpg')
        srcLabel = os.path.join(txt_dir, name + '.txt')
 
        if i in train:
            dst_train_Image = os.path.join(img_train_path, name + '.jpg')
            dst_train_Label = os.path.join(label_train_path, name + '.txt')
            shutil.copyfile(srcImage, dst_train_Image)
            shutil.copyfile(srcLabel, dst_train_Label)
        elif i in val:
            dst_val_Image = os.path.join(img_val_path, name + '.jpg')
            dst_val_Label = os.path.join(label_val_path, name + '.txt')
            shutil.copyfile(srcImage, dst_val_Image)
            shutil.copyfile(srcLabel, dst_val_Label)

 

if __name__ == '__main__':
    """
    python split_datasets.py --image-dir my_datasets/color_rings/imgs --txt-dir my_datasets/color_rings/txts --save-dir my_datasets/color_rings/train_data
    """
    parser = argparse.ArgumentParser(description='split datasets to train,val,test params')
    parser.add_argument('--image-dir', type=str,default='datasets/obb/JPEGImages', help='image path dir')
    parser.add_argument('--txt-dir', type=str,default='datasets/obb/labelTxt' , help='txt path dir')
    parser.add_argument('--save-dir', default='datasets/obb/meter', type=str, help='save dir')
    args = parser.parse_args()
    image_dir = args.image_dir
    txt_dir = args.txt_dir
    save_dir = args.save_dir
 
    main(image_dir, txt_dir, save_dir)

运行完后,生成一个文件夹(meter),里面是划分好的数据集。

4、将labels的data格式转换成yolo格式

在项目根目录下新建一个data_to_yolo.py文件夹。

import sys
from ultralytics.data.converter import convert_dota_to_yolo_obb

# sys.path.append(r'/ultralytics')

convert_dota_to_yolo_obb(r'datasets/obb/meter')

注意修改路径!!

!!!!运行前需修改两个位置!!!!

ctrl+左键定位 这个函数 convert_dota_to_yolo_obb 直接可以定位到 converter.py 文件中

**将class_mapping修改成自己的类名**

**修改图片后缀**

然后运行data_to_yolo.py,运行完后在labels文件夹下会生成train和val两个文件夹。

至此,数据集(meter)完成,下一步可以用这个数据集训练

三、训练

3.1 在项目根目录下新建一个obb.yaml文件,复制以下内容。注意修改路径

names:
  0: label
path: E:\python\yolov8\ultralytics_train_7_30\datasets/obb/meter
train: images/train
val: images/val

3.2 在ultralytics/models/yolo/obb下新建一个obb_train_my.py文件。复制以下内容

(需要提前下载yolov8n.pt 和yolov8n-obb.pt两个模型在这个目录下)前往官网下载

from ultralytics import YOLO

# Load a model
model = YOLO("yolov8n-obb.pt")  # load a pretrained model (recommended for training)


# Train the model
results = model.train(task='obb',data="obb.yaml", model='yolov8n-obb.pt', epochs=300, imgsz=1024, batch=16, workers=0,
                      resume=False)

3.3 运行obb_train_my.py,开始训练

根据自己的情况修改参数,或者前往ultralytics/cfg/default.yaml,这个路径下修改这个yaml文件里面的参数,都是一样的

注意,参数设置优先是results = model.train(task='detect',data="datasets/detect/dataset.yaml", model='yolov8n.pt', epochs=300, imgsz=640, batch=16, workers=0,                       resume=False)这个里面的。

如果采用default.yaml,则需要将上面代码中的参数删除

eg:results = model.train()

3.2 

训练后结果会保存在ultralytics/models/y/obb/runs下面

  • 16
    点赞
  • 11
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值