1、创建目录如下所示:
VOC2012
---Annotations
---ImageSets
---Main
---JPEGImages
2、标定图片放到文件夹VOC2012/JPEGImages,使用labelImg工具进行标定,将标定后的xml文件剪切到文件夹VOC2012/Annotations
3、脚本修改xml文件中红色方框内容
4、脚本在VOC2012/JPEGImages/Main文件夹生成:类别名_train.txt和类别名_val.txt;同时生成trainval.txt、train.txt、val.test、test.txt
脚本程序如下所示,脚本放到文件夹VOC2012运行
import xml.etree.ElementTree as ET
import os
import random
import cv2 as cv
def change_image_format(old_format='.png', new_format='.jpg'):
img_dir = "../VOC2012/JPEGImages/"
files = os.listdir(img_dir)
for img_file in files:
if os.path.isfile(os.path.join(img_dir, img_file)):
image_path = os.path.join(img_dir, img_file)
# print(image_path)
image = cv.imread(image_path)
new_image_path = image_path.replace(old_format, new_format)
cv.imwrite(new_image_path, image, [cv.IMWRITE_JPEG_QUALITY, 100])
print("processed image : %s" % (new_image_path))
def xml_modification():
ann_dir = "../VOC2012/Annotations/"
img_dir = "C:\\Users\\mengkun\\Desktop\\tmp\\VOC2012\\JPEGImages\\" # 改成自己数据集JPEGImages文件夹绝对路径
files = os.listdir(ann_dir)
for xml_file in files:
if os.path.isfile(os.path.join(ann_dir, xml_file)):
xml_path = os.path.join(ann_dir, xml_file)
# print(xml_path)
tree = ET.parse(xml_path)
root = tree.getroot()
for elem in root.iter('folder'):
elem.text = 'voc2012'
for elem in root.iter('filename'):
pass
for elem in root.iter('path'):
path = elem.text
filename = path.split('\\')[-1]
new_path = img_dir + filename
elem.text = new_path
tree.write(xml_path)
print("processed xml : %s" % (xml_path))
def generate_classes_txt():
ann_dir = "../VOC2012/Annotations/"
ok_train = open("../VOC2012/ImageSets/Main/ok_train.txt", 'w') # 改成自己数据集类别
ok_val = open("../VOC2012/ImageSets/Main/ok_val.txt", 'w') # 改成自己数据集类别
spot_train = open("../VOC2012/ImageSets/Main/spot_train.txt", 'w') # 改成自己数据集类别
spot_val = open("../VOC2012/ImageSets/Main/spot_val.txt", 'w') # 改成自己数据集类别
files = os.listdir(ann_dir)
for xml_file in files:
if os.path.isfile(os.path.join(ann_dir, xml_file)):
xml_path = os.path.join(ann_dir, xml_file)
# print(xml_path)
tree = ET.parse(xml_path)
root = tree.getroot()
for elem in root.iter('filename'):
filename = elem.text
name_list = []
for elem in root.iter('name'):
name = elem.text
name_list.append(name)
if "ok" in name_list:
ok_train.write(filename.replace(".jpg", " ") + str(1) + "\n")
ok_val.write(filename.replace(".jpg", " ") + str(1) + "\n")
else:
ok_train.write(filename.replace(".jpg", " ") + str(-1) + "\n")
ok_val.write(filename.replace(".jpg", " ") + str(-1) + "\n")
if "spot" in name_list:
spot_train.write(filename.replace(".jpg", " ") + str(1) + "\n")
spot_val.write(filename.replace(".jpg", " ") + str(1) + "\n")
else:
spot_train.write(filename.replace(".jpg", " ") + str(-1) + "\n")
spot_val.write(filename.replace(".jpg", " ") + str(-1) + "\n")
ok_train.close()
ok_val.close()
spot_train.close()
spot_val.close()
def generate_train_val_test_txt():
xml_file_path = "../VOC2012/Annotations/"
save_Path = "../VOC2012/ImageSets/Main/"
trainval_percent = 0.9
train_percent = 0.9
total_xml = os.listdir(xml_file_path)
num = len(total_xml)
list = range(num)
tv = int(num * trainval_percent)
tr = int(tv * train_percent)
trainval = random.sample(list, tv)
train = random.sample(trainval, tr)
print("train and val size", tv)
print("train size", tr)
ftrainval = open(os.path.join(save_Path, 'trainval.txt'), 'w')
ftest = open(os.path.join(save_Path, 'test.txt'), 'w')
ftrain = open(os.path.join(save_Path, 'train.txt'), 'w')
fval = open(os.path.join(save_Path, 'val.txt'), 'w')
for i in list:
name = total_xml[i][:-4] + '\n'
if i in trainval:
ftrainval.write(name)
if i in train:
ftrain.write(name)
else:
fval.write(name)
else:
ftest.write(name)
ftrainval.close()
ftrain.close()
fval.close()
ftest.close()
xml_modification()
# generate_classes_txt()
generate_train_val_test_txt()