（25）运动目标检测之旋转目标框 xml 转YOLO格式及数据集划分
- 针对旋转目标框： xml 转 txt
- 数据集{训练集与验证集} 划分
- 矩形框xml 转旋转矩形框xml
xml转txt
python
复制代码
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)
数据集划分
python
复制代码
# -*- 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
python
复制代码
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()