这个程序可以从数据集中随机抽取训练集和测试集
原来的数据结构为:
data 类别1 image1 image2 image3 …… 类别2 image1 image2 image3 …… ……
抽取完的数据结构为:
data 类别1 image3 …… 类别2 image2 image3 …… …… test 类别1 image3 …… 类别2 image1 …… ……
然后将data文件夹改为train就得到了训练集和测试集,如果还需要验证集就接着再抽取一次,(对了,抽取之前建立一个test文件夹)
import os import random import shutil # source_file:源路径, target_ir:目标路径 def cover_files(source_dir, target_ir): for file in os.listdir(source_dir): source_file = os.path.join(source_dir, file) if os.path.isfile(source_file): shutil.copy(source_file, target_ir) def ensure_dir_exists(dir_name): """Makes sure the folder exists on disk. Args: dir_name: Path string to the folder we want to create. """ if not os.path.exists(dir_name): os.makedirs(dir_name) def moveFile(file_dir, save_dir): ensure_dir_exists(save_dir) path_dir = os.listdir(file_dir) #取图片的原始路径 filenumber=len(path_dir) rate=0.1667 #自定义抽取图片的比例,比方说100张抽10张,那就是0.1 picknumber=int(filenumber*rate) #按照rate比例从文件夹中取一定数量图片 print(picknumber) sample = random.sample(path_dir, picknumber) #随机选取picknumber数量的样本图片 # print (sample) for name in sample: shutil.move(file_dir+name, save_dir+name) def mkdir(path): folder = os.path.exists(path) if not folder: #判断是否存在文件夹如果不存在则创建为文件夹 os.makedirs(path) #makedirs 创建文件时如果路径不存在会创建这个路径 print("--- new folder... ---") print("--- OK ---") else: print("--- There is this folder! ---") if __name__ == '__main__': path='/home/megstudio/workspace/image/' dirs = os.listdir( path+'imageall/' ) for file in dirs: file_dir = path+'imageall/'+file+'/' #源图片文件夹路径 print(file_dir) save_dir = path+'test/'+file #移动到新的文件夹路径 print(save_dir) mkdir(save_dir) #创造文件夹 save_dir = save_dir+'/' moveFile(file_dir, save_dir)
需要再用pytorch导入模型就参考我的下一篇博文。
pytorch导入自建数据集