import xml.etree.ElementTree as ET
import os
import argparse
import random
import shutil
rootpath = os.getcwd()
classes = ["0", "1"]
class VocToYolo(object):
'''
labelimg贴的标签voc文件转化为yolo支持的格式
'''
def __init__(self,xml_path,save_txt_path):
self.xml_path =xml_path
self.save_txt_path = save_txt_path
self.roots = self.readxmls()
def readxmls(self):
f = open(self.xml_path)
xml_text = f.read()
root = ET.fromstring(xml_text)
f.close()
return root
def get_image_path(self):
root = self.readxmls()
images_path = root.find('path')
img_name = images_path.text.split('/')[-1].split('.')[0]
print(img_name)
return images_path.text,img_name
def get_hw(self):
size = self.roots.find('size')
w = int(size.find('width').text)
h = int(size.find('height').text)
return (w,h)
def convert_size(self,size, box):
dw = 1.0 / size[0]
dh = 1.0 / size[1]
x = (box[0] + box[1]) / 2.0
y = (box[2] + box[3]) / 2.0
w = box[1] - box[0]
h = box[3] - box[2]
x = x * dw
w = w * dw
y = y * dh
h = h * dh
return (x, y, w, h)
def get_xyxy_label(self):
_,a = self.get_image_path()
os.makedirs(self.save_txt_path,exist_ok=True)
out_file = open(self.save_txt_path+'/' + str(a) + '.txt', 'w')
for obj in self.roots.iter('object'):
cls = obj.find('name').text
if cls not in classes:
print(cls)
continue
print("标签:",cls)
label_now = classes.index(cls)
print("标签索引:", label_now)
xmlbox = obj.find('bndbox')
xmin = int(float(xmlbox.find('xmin').text))
xmax = int(float(xmlbox.find('ymin').text))
ymin = int(float(xmlbox.find('xmax').text))
ymax = int(float(xmlbox.find('ymax').text))
size = self.get_hw()
xyhw = self.convert_size(size, (xmin, ymin, xmax, ymax))
out_file.write(str(label_now) + " " + " ".join([str(xyhw) for xyhw in xyhw]) + '\n')
print(xmin, ymin, xmax, ymax)
return 0
def excute_voc_to_yolo():
'''
转化为yolo格式主函数
'''
parser = argparse.ArgumentParser()
parser.add_argument('--xml_path', type=str, default='./Annotations', help='xml路径')
parser.add_argument('--save_txt', type=str, default='./labels', help='txt保存路径')
opt = parser.parse_args()
for i in os.listdir(opt.xml_path):
xml_path = opt.xml_path + '/' + str(i)
ab = VocToYolo(xml_path, opt.save_txt)
ab.get_xyxy_label()
class SplitValTrain(object):
'''
划分验证集和训练集
'''
def __init__(self,rootpath,label_path,images_path,ratios):
self.rootpath = rootpath
self.label_path = label_path
self.images_path = images_path
self.ratios = ratios
self.train_images_path = rootpath + '/yolodata/train/images'
self.train_labels_path = rootpath + '/yolodata/train/labels'
self.val_images_path = rootpath + '/yolodata/val/images'
self.val_labels_path = rootpath + '/yolodata/val/labels'
def make_data_dir(self):
os.makedirs(self.train_images_path,exist_ok=True)
os.makedirs(self.train_labels_path,exist_ok=True)
os.makedirs(self.val_images_path, exist_ok=True)
os.makedirs(self.val_labels_path, exist_ok=True)
def split_val_train(self):
self.make_data_dir()
filename = os.listdir(self.images_path)
# print(filename)
random.shuffle(filename)
go_to = int(self.ratios*len(filename))
print(go_to,len(filename))
for i in range(len(filename)):
name_num = filename[i].split('.')[0]
print(filename[i])
if i<go_to:
shutil.copy(self.images_path+'/'+str(filename[i]), self.train_images_path)
shutil.copy(self.label_path + '/' + str(name_num) + '.txt', self.train_labels_path)
else:
shutil.copy(self.images_path + '/' + str(filename[i]), self.val_images_path)
shutil.copy(self.label_path + '/' + str(name_num) + '.txt', self.val_labels_path)
return 0
def excute_split_val_train():
parser = argparse.ArgumentParser()
parser.add_argument('--rootpath', type=str, default=rootpath, help='当前目录')
parser.add_argument('--label_path', type=str, default='./labels', help='txt文件路径')
parser.add_argument('--images_path', type=str, default='./JPEGImages', help='图像路径')
parser.add_argument('--ratios', type=float, default=0.8, help='训练集划分比例')
opt = parser.parse_args()
ab = SplitValTrain(opt.rootpath, opt.label_path, opt.images_path, opt.ratios)
ab.split_val_train()
if __name__ == '__main__':
# excute_voc_to_yolo()
excute_split_val_train()
12、voc数据集转yolo格式、划分yolo训练集和验证集
最新推荐文章于 2024-07-09 17:20:39 发布