需求:
现在的数据集形式:
举例: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