需求:
现在的数据集形式:
举例:train文件夹下的形式
也即现在图像数据已经按照标签划分到了相应的文件夹下,然后想把所有的图像数据统一放到’./images’文件夹下,同时生成相应的索引文件,方便后续适配自己继承实现的torch.utils.data.Dataset类。
代码:
import os
import shutil
import re
import numpy as np
import pandas as pd
train_source = './train'
val_source = './val'
test_source = './test'
target = './images/'
def Process_images(rootdir, setname='train'):
print(setname)
file_list = os.listdir(rootdir)
image_name = []
image_label = []
for i in range(len(file_list)):
label_name = file_list[i]
path_name = os.path.join(rootdir, label_name)
print("path_name {} ".format(path_name))
imgfile = os.listdir(path_name)
for img in imgfile:
image_name.append(img)
image_label.append(label_name)
img_path = os.path.join(path_name, img)
print("img_path {} ".format(img_path))
shutil.move(img_path, target)
save_all = {"filename": image_name, "label": image_label}
save_data = pd.DataFrame(save_all)
save_name = setname + '.csv'
save_data.to_csv(save_name, index=0) # 取消索引
Process_images(train_source, setname='train')
Process_images(val_source, setname='val')
Process_images(test_source, setname='test')
结果示意:
test.csv