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

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

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

 

上一篇:用Python爬虫下载漫画


下一篇:谷粒商城学习——P90 调用远程服务