完整数据集中随机抽取训练集和测试集

这个程序可以从数据集中随机抽取训练集和测试集
原来的数据结构为：

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导入自建数据集