一、图片数据文件切分
代码实现规则:
1、首先要有一个未划分前的数据文件夹,此文件夹内包含多个子文件夹,每个子文件夹包含不同类别的图片。
2、要为划分后数据建立一个的空的存储文件夹,下面的代码会在空文件夹中自动生成“train”“test”文件夹,并将按比例划分好的图片数据存储进去。
存储后的datas_splited文件夹长这样:
import os
from shutil import copy, rmtree
import random
from tqdm import tqdm
def main():
'''
split_rate : 测试集划分比例
init_dataset: 未划分前的数据集路径
new_dataset : 划分后的数据集路径
'''
def makedir(path):
if os.path.exists(path):
rmtree(path)
os.makedirs(path)
split_rate = 0.2
init_dataset = r'datas\BAT'
new_dataset = r'datasets'
random.seed(0)
classes_name = [name for name in os.listdir(init_dataset)]
makedir(new_dataset)
training_set = os.path.join(new_dataset, "train")
test_set = os.path.join(new_dataset, "test")
makedir(training_set)
makedir(test_set)
for cla in classes_name:
makedir(os.path.join(training_set, cla))
makedir(os.path.join(test_set, cla))
for cla in classes_name:
class_path = os.path.join(init_dataset, cla)
img_set = os.listdir(class_path)
num = len(img_set)
test_set_index = random.sample(img_set, k=int(num*split_rate))
with tqdm(total=num,desc=f'Class : ' + cla, mininterval=0.3) as pbar:
for _, img in enumerate(img_set):
if img in test_set_index:
init_img = os.path.join(class_path, img)
new_img = os.path.join(test_set, cla)
copy(init_img, new_img)
else:
init_img = os.path.join(class_path, img)
new_img = os.path.join(training_set, cla)
copy(init_img, new_img)
pbar.update(1)
print()
if __name__ == '__main__':
main()