（25）运动目标检测之旋转目标框 xml 转YOLO格式及数据集划分

（25）运动目标检测之旋转目标框 xml 转YOLO格式及数据集划分

• 针对旋转目标框： xml 转 txt
• 数据集{训练集与验证集} 划分
• 矩形框xml 转旋转矩形框xml

xml转txt

import os
import math
import xml.etree.ElementTree as ET
import numpy as np
import cv2


def rotate(cx, cy, w, h, angle):
    # 用于常规坐标时，angle是顺时针旋转角度
    # 用于图像坐标时，angle是逆时针旋转角度
    # math.cos(angle)这里的angle单位是rad
    # angle = angle/180*math.pi  单位转换
    angle = -angle  # 这里是图像坐标
    points = [[cx - w / 2, cy - h / 2], [cx + w / 2, cy - h / 2],
              [cx + w / 2, cy + h / 2], [cx - w / 2, cy + h / 2]]
    newpoints = []
    if angle < 0:  # 逆时针
        angle = -angle
        for point in points:
            x, y = point
            newx = round((x - cx) * math.cos(angle) - (y - cy) * math.sin(angle) + cx, 1)
            newy = round((x - cx) * math.sin(angle) + (y - cy) * math.cos(angle) + cy, 1)
            newpoints.append([newx, newy])
    else:
        for point in points:
            x, y = point
            newx = round((x - cx) * math.cos(angle) + (y - cy) * math.sin(angle) + cx, 1)
            newy = round((y - cy) * math.cos(angle) - (x - cx) * math.sin(angle) + cy, 1)
            newpoints.append([newx, newy])
    return newpoints


def roxml2txt(xml_files_path, save_txt_files_path, classes):
    # dir是xml文件目录
    files = os.listdir(xml_files_path)
    for f in files:
        xml_file_ = os.path.join(xml_files_path, f)
        out_txt_path = os.path.join(save_txt_files_path, f.split('.')[0] + '.txt')
        img = cv2.imread(os.path.join('images', f.split('.')[0] + '.jpg'))
        xml_file = open(xml_file_, encoding='gbk')
        root = ET.parse(xml_file).getroot()
        # xml = ET.parse(os.path.join(dir,f))
        # root = xml.getroot()
        size = root.find('size')
        img_w = int(size.find('width').text)
        img_h = int(size.find('height').text)
        boxes = root.iter('robndbox')
        with open(out_txt_path, 'w+') as t:
            for obj in root.iter('object'):
                difficult = obj.find('difficult').text
                cls = obj.find('name').text
                if cls not in classes or int(difficult) == 1:
                    continue
                cls_id = classes.index(cls)
                box = obj.find('robndbox')

                cx = float(box.find('cx').text)
                cy = float(box.find('cy').text)
                w = float(box.find('w').text)
                h = float(box.find('h').text)
                angle = float(box.find('angle').text)
                newpoints = rotate(cx, cy, w, h, angle)  # 计算旋转后的4个点坐标
                # 用于查看坐标转换是否正确，在原图上画矩形框，自行修改图片路径
                newpoints = np.array(newpoints)
                newpoints = newpoints.astype(int)

                cv2.polylines(img, [newpoints], isClosed=True, color=(0, 0, 255))

                ##########################################
                line = str(cls_id)+' '
                for point in newpoints:
                    line += str(point[0]/img_w) + ' ' + str(point[1]/img_h) + ' '
                t.write(line+"\n")
                print(line)
            t.close()
        cv2.namedWindow("pic",0)
        cv2.imshow("pic", img)
        cv2.waitKey(10)


if __name__ == "__main__":

    classes = ['corner',  'c1']  # 改成自己的类别

    xml_files1 = r'./xmls'
    # 2、转化为yolo格式的txt标签文件存储路径
    save_txt_files1 = r'./labels'

    roxml2txt(xml_files1, save_txt_files1, classes)

数据集划分

# -*- coding:utf-8 -*
import os
import random


"""
time:20220729
writer:yohn
function:将数据集按照 比例随机 分为训练集、测试集和验证集,在此就将测试与验证归为一个了
"""

import os
import shutil


def data_split(full_list, ratio):
    """
    数据集拆分: 将列表full_list按比例ratio（随机）划分为2个子列表sublist_1与sublist_2
    :param full_list: 数据列表
    :param ratio:     子列表1
    """
    n_total = len(full_list)
    offset = int(n_total * ratio)
    if n_total == 0 or offset < 1:
        return [], full_list

    random.shuffle(full_list)
    sublist_1 = full_list[:offset]
    sublist_2 = full_list[offset:]
    return sublist_1, sublist_2


train_p="./train"
val_p="./val"
imgs_p="images\\src1"
labels_p="labels\\src1"

#创建训练集
if not os.path.exists(train_p):#指定要创建的目录
    os.mkdir(train_p)
tp1=os.path.join(train_p,imgs_p)
tp2=os.path.join(train_p,labels_p)
print(tp1,tp2)
if not os.path.exists(tp1):#指定要创建的目录
    os.mkdir(tp1)
if not os.path.exists(tp2):  # 指定要创建的目录
    os.mkdir(tp2)

#创建测试集文件夹
if not os.path.exists(val_p):#指定要创建的目录
    os.mkdir(val_p)
vp1=os.path.join(val_p,imgs_p)
vp2=os.path.join(val_p,labels_p)
print(vp1,vp2)
if not os.path.exists(vp1):#指定要创建的目录
    os.mkdir(vp1)
if not os.path.exists(vp2):  # 指定要创建的目录
    os.mkdir(vp2)

#数据集路径
path1=".\\images"
path2=".\\labels"
#划分数据集，设置数据集数量占比
proportion_ = 1 #训练集占比

total_file = os.listdir(path1)

num = len(total_file)  # 统计所有的标注文件
list_=[]
for i in range(0,num):
    list_.append(i)

list1,list2=data_split(list_,proportion_)

for i in range(0,num):
    file=total_file[i]
    print(i,' - ',total_file[i])
    name=file.split('.')[0]
    if i in list1:
        jpg_1 = os.path.join(path1, file)
        jpg_2 = os.path.join(train_p, imgs_p, file)
        txt_1 = os.path.join(path2, name + '.txt')
        txt_2 = os.path.join(train_p, labels_p, name + '.txt')
        shutil.copyfile(jpg_1, jpg_2)
        shutil.copyfile(txt_1, txt_2)

    elif i in list2:
        jpg_1 = os.path.join(path1, file)
        jpg_2 = os.path.join(val_p, imgs_p, file)
        txt_1 = os.path.join(path2, name + '.txt')
        txt_2 = os.path.join(val_p, labels_p, name + '.txt')
        shutil.copyfile(jpg_1, jpg_2)
        shutil.copyfile(txt_1, txt_2)

print("数据集划分完成： 总数量：",num," 训练集数量：",len(list1)," 验证集数量：",len(list2))

矩形框xml 转旋转矩形框xml

def RectxmltoRoRectxml(path,xmlname, savexmlpath):
    '''
    :param srcxmlpath:  输入xml路径
    :param dstxmlpath:  保存的xml路径，如果本地不存在文件夹，将重新创建 ../convert
    :return:  NULl
    '''

    #定义将要保存的xml结果文件
    node_root = Element('annotation')
    node_folder = SubElement(node_root, 'folder')
    node_folder.text = 'Images'
    node_filename = SubElement(node_root, 'filename')
    node_size = SubElement(node_root, 'size')
    node_width = SubElement(node_size, 'width')
    node_height = SubElement(node_size, 'height')
    node_depth = SubElement(node_size, 'depth')

    #读取通过labelimg标注好的xml文件
    xmlpath=os.path.join(path,xmlname)
    with open(xmlpath, 'r',encoding='utf-8') as f:
        tree = ET.parse(f)
    root = tree.getroot()
    # print(root)
    # print(list(root))

    node_filename.text = root.find("filename").text
    size = root.find("size")
    node_width.text = size.find("width").text
    node_height.text = size.find("height").text
    node_depth.text = size.find("depth").text
    # print("w = ",w.text,"h = ",h.text)

    # 查看多个重复元素
    for obj in root.iter("object"):
        box = obj.find("bndbox")
        x_min = int(box.find("xmin").text)
        y_min = int(box.find("ymin").text)
        x_max = int(box.find("xmax").text)
        y_max = int(box.find("ymax").text)


        #往新的xml中赋值
        node_object = SubElement(node_root, 'object')
        node_type = SubElement(node_object, 'type')
        node_type.text="robndbox"
        node_name = SubElement(node_object, 'name')
        node_name.text=obj.find("name").text
        node_difficult = SubElement(node_object, 'difficult')
        node_difficult.text = '0'
        node_bndbox = SubElement(node_object, 'robndbox')
        node_xmin = SubElement(node_bndbox, 'cx')
        node_xmin.text = str(round(0.5*(x_min+x_max)))
        node_ymin = SubElement(node_bndbox, 'cy')
        node_ymin.text = str(round(0.5*(y_min+y_max)))
        node_w = SubElement(node_bndbox, 'w')
        node_w.text = str(x_max-x_min)
        node_h = SubElement(node_bndbox, 'h')
        node_h.text = str(y_max-y_min)
        node_angle = SubElement(node_bndbox, 'angle')
        node_angle.text = "0"

    xml = tostring(node_root, pretty_print=True)  # 格式化显示，该换行的换行

    filename = os.path.join(savexmlpath,xmlname)
    f = open(filename, "wb")
    f.write(xml)
    f.close()