yolov8-obb旋转目标检测详细流程

news2024/12/27 11:49:15

一、数据标注

由于旋转目标标注与传统的目标检测标注有一定的区别，故使用roLabelImg或者labelimg2进行标注，以下标注软件选用的是roLabelImg。
roLabelImg下载地址：https://github.com/cgvict/roLabelImg
直接在原有的labelimg环境中使用roLabelImg,运行python roLabelImg.py
在这里插入图片描述
打开标注软件后，e表示标注旋转目标，w表示标注正常目标，d表示下一张，a表示上一张，c表示小幅度顺时针旋转，x表示小幅度逆时针旋转，v表示大幅度顺时针旋转，z表示大幅度逆时针旋转。

二、数据集处理

整体处理流程xml —> dota_xml —> dota_txt ----> txt
文件夹
在这里插入图片描述
第一步：xml —> 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 = ['1', '2', '3', '4']  #修改为自己的标签
 
 
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 = '/home/y/code/yolov8/data/lase2/labels/val_xml'
    dotaxml_path = '/home/y/code/yolov8/data/lase2/labels/val_dota'
    out_path = '/home/y/code/yolov8/data/lase2/labels/val_original'
    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)

第二步：dota_txt —> txt

import sys
import inspect
sys.path.insert(0, '/home/y/code/yolov8')
print(sys.path)
 
from ultralytics.data.converter import convert_dota_to_yolo_obb
source_file = inspect.getsourcefile(convert_dota_to_yolo_obb)
print("convert_dota_to_yolo_obb 函数所在文件的路径：", source_file)
convert_dota_to_yolo_obb('/home/y/code/yolov8/data/lase2')
#关于dataobb文件下的目录下面会详细说明

注：在进行第二步之前需要做的一些代码修改
1、修改yolov8/ultralytics/data/converter.py
eg：在这里插入图片描述
将class_mapping修改为自己的标签。
2、修改yolov8/ultralytics/cfg/models/v8/yolov8-obb.yaml

将nc修改为自己的类别数量。
3、创建训练所需的yaml文件放置在数据集文件夹中

# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
path: /home/y/code/yolov8/data/lase2
train: images/train  #  images
val: images/val  #  images

# number of classes
nc: 4

# class names
names: 
  0: ture
  1: ban
  2: man
  3: di

三、开始训练

1、更改/home/wangyu/code/yolov8/ultralytics/cfg/default.yaml
在这里插入图片描述
2、开始训练

import argparse
import sys
sys.path.append('/home/y/code/yolov8/ultralytics')
from ultralytics import YOLO


def parse_opt(known=False):
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', type=str, default='yolov8s-obb.pt', help='initial weights path')
    parser.add_argument('--epochs', type=int, default=1600, help='total training epochs')
    parser.add_argument('--imgsz', type=int, default=640, help='train, val image size (pixels)')
    parser.add_argument('--batch', type=str, default=2, help='total batch size for all GPUs, -1 for autobatch')
    parser.add_argument('--lr0', type=str, default=0.01, help=' (float) initial learning rate (i.e. SGD=1E-2, Adam=1E-3)')
    parser.add_argument('--cls', type=str, default=1.5, help=' (float) cls loss gain (scale with pixels)')
    parser.add_argument('--data', type=str, default='/home/y/code/yolov8/data/lase/my-dota8-obb.yaml', help='dataset.yaml path')    parser.add_argument('--workers', type=str, default=8)
    parser.add_argument('--device', type=str, default=0, help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    return parser.parse_known_args()[0] if known else parser.parse_args()

def main(opt):
    model = YOLO(opt.model)
    model.train(data=opt.data, epochs=opt.epochs, imgsz=opt.imgsz, batch=opt.batch, device=opt.device, lr0=opt.lr0, cls=opt.cls)
    model.val(imgsz=opt.imgsz)


if __name__ == '__main__':
    # results,metrics,success = main()
    opt = parse_opt()
    main(opt)