88 lines
3.4 KiB
Python
Executable File
88 lines
3.4 KiB
Python
Executable File
# -*- coding: utf-8 -*-
|
|
"""
|
|
author: Hatcher W
|
|
use:
|
|
拆分训练集和测试集，并在ImageSet/Main中生成train.txt和test.txt
|
|
生成yolo使用的train_img_path.txt和test_img_path.txt文件
|
|
可以仅抽选一定比例
|
|
"""
|
|
|
|
import os
|
|
import xml.etree.ElementTree as ET
|
|
import argparse
|
|
|
|
|
|
def split_train_test_set(root_dir, target_dir, test_fraction=0.1):
|
|
"""
|
|
split dataset into test and train subsets. The script will create train.txt and test.txt
|
|
that contain a list of files for VOC dataset, and train_img_path.txt/test_img_path.txt
|
|
for yolo training. All the four files will be created in ImageSet/Main directory.
|
|
:param test_fraction: the fraction of test set
|
|
:return: None
|
|
"""
|
|
train_txt = open(os.path.join(target_dir, "train.txt"), "w")
|
|
test_txt = open(os.path.join(target_dir, "test.txt"), "w")
|
|
yolo_train_img_path = open(os.path.join(target_dir, 'train_img_path.txt'), 'w')
|
|
yolo_test_img_path = open(os.path.join(target_dir, 'test_img_path.txt'), 'w')
|
|
xml_path = os.path.join(root_dir, 'Annotations')
|
|
'''从已有标注数据中抽选'''
|
|
select_percent = 1
|
|
# 测试集合占比
|
|
# test_percent = 0.01
|
|
xml_num = len(os.listdir(xml_path))
|
|
print('=========based on labels(xml): {} ============'.format(xml_num))
|
|
|
|
train_num = 0
|
|
test_num = 0
|
|
total_processed_num = 0
|
|
if test_fraction <= 1e-10:
|
|
sample_interval = xml_num + 1
|
|
else:
|
|
sample_interval = int(1 / test_fraction)
|
|
select_interval = int(1 / select_percent)
|
|
print('select sample interval is: {}'.format(select_interval))
|
|
print('test sample interval is: {}'.format(sample_interval))
|
|
|
|
files = os.listdir(xml_path)
|
|
for file in files:
|
|
total_processed_num = total_processed_num + 1
|
|
if total_processed_num % select_interval != 0:
|
|
if total_processed_num % 100 == 0:
|
|
print('completed {} / {}'.format(total_processed_num, xml_num))
|
|
continue
|
|
|
|
img_name = file.split('.')[0]
|
|
tree = ET.parse(os.path.join(xml_path, file))
|
|
root = tree.getroot()
|
|
file_name = root.find('filename').text
|
|
|
|
img_path = os.path.join(root_dir, 'JPEGImages', file_name)
|
|
|
|
if total_processed_num % (sample_interval * select_interval) == 0:
|
|
test_txt.write(str(img_name) + '\n')
|
|
yolo_test_img_path.write(img_path + '\n')
|
|
test_num = test_num + 1
|
|
elif total_processed_num % select_interval == 0:
|
|
train_txt.write(str(img_name) + '\n')
|
|
yolo_train_img_path.write(img_path + '\n')
|
|
train_num = train_num + 1
|
|
if total_processed_num % 100 == 0:
|
|
print('completed {} / {}'.format(total_processed_num, xml_num))
|
|
print('total:', total_processed_num, 'train:', train_num, 'test:', test_num)
|
|
train_txt.close()
|
|
test_txt.close()
|
|
yolo_train_img_path.close()
|
|
yolo_test_img_path.close()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
root_dir = "../data_MPKFNet/data/VOC/VOCdevkit/VOC16/"
|
|
parser = argparse.ArgumentParser(description="define all the file paths")
|
|
parser.add_argument("--target_dir", type=str, help="image directory",
|
|
default=os.path.join(root_dir, 'ImageSets', 'Main'))
|
|
args = parser.parse_args()
|
|
target_dir = args.target_dir
|
|
if not os.path.exists(target_dir):
|
|
os.makedirs(target_dir)
|
|
|
|
split_train_test_set(root_dir, target_dir, 0.0) |