Python 处理图像分类数据集之训练集、测试集分离(随机)

本文代码可推广到其他数据集上，也可以推广到您想划分训练集、验证集、测试集上，下面以自己的数据集和划分训练集、测试集为例！
本段代码的功能是以一定的比例来划分数据集为测试集、训练集，且在不改变原数据集的情况下把属于每一个类别的训练集和测试集图片copy到另外两个文件夹train、test下：

只需给出上图文件夹所在绝对路径、训练集和测试集的划分比例即可:

import os
import random
import shutil
import csv
import numpy as np
def CopyFile(imageDir,test_rate,save_test_dir,save_train_dir):#三个参数，第一个为每个类别的所有图像在计算机中的位置
    #第二个为copy的图片数目所占总的比例，最后一个为移动的图片保存的位置，
    image_number = len(imageDir)  #图片总数目
    test_number = int(image_number * test_rate)#要移动的图片数目
    print("要移动到%s目录下的图片数目为:%d"%(save_test_dir,test_number))
    test_samples = random.sample(imageDir, test_number)#随机截取列表imageDir中数目为test_number的元素
 # copy图像到目标文件夹
    if not os.path.exists(save_test_dir):
        os.makedirs(save_test_dir)
        print("save_test_dir has been created successfully!")
    else:
        print("save_test_dir already exited!")
    if not os.path.exists(save_train_dir):
        os.makedirs(save_train_dir)
        print("save_train_dir has been created successfully!")
    else:
        print("save_train_dir already exited!")
    for i,j in enumerate(test_samples):
        shutil.copy(test_samples[i], save_test_dir+test_samples[i].split("/")[-1])
    print("tets移动完成！")
    for train_imgs in imageDir:
        if train_imgs not in test_samples:
            shutil.copy(train_imgs, save_train_dir+train_imgs.split("/")[-1])
    print("train移动完成")
#只需给定file_path、test_rate即可完成整个任务
#原始路径+分割比例
################################
file_path="D:\\数据\\images"
test_rate = 0.5
################################
file_dirs=os.listdir(file_path)
origion_paths=[]
save_test_dirs=[]
save_train_dirs=[]
for path in file_dirs:
   origion_paths.append(file_path+"/"+path+"/")
   save_train_dirs.append("./train/"+path+"/")
   save_test_dirs.append("./test/"+path+"/")
for i,origion_path in enumerate(origion_paths):
    image_list = os.listdir(origion_path) #获得原始路径下的所有图片的name（默认路径下都是图片）
    image_Dir=[]
    for x,y in enumerate(image_list):
        image_Dir.append (os.path.join(origion_path, y))
    print("%s目录下共有%d张图片！"%(origion_path,len(image_Dir)))
    CopyFile(image_Dir,test_rate,save_test_dirs[i],save_train_dirs[i])
print("all datas has been moved successfully!")

结果:



为什么这样做呢？因为有时候跑论文用到ImageFolder、DataLoader类，会直接从train、test文件夹下读取图片！是为了满足数据集能被直接加载的要求！
详细可加：https://blog.csdn.net/hb_learing/article/details/115455269在这篇论文复现过程中，就对数据集的格式提出了要求！哈哈，CUB200数据集也是同样的操作方法，可以用CUB200数据集检验哦。

因为是按照1:1的比例划分的，我怕测试集的数目比训练集还多，加入了如下的代码来统计下训练集、测试集、总的数据集的数目：

测试集:

#统计测试集数目
import os
path="D:\\数据\\216_2021\\test"
sum_test=0
dirs=os.listdir(path)
for Dir in dirs:
    new_dir=os.path.join(path,Dir)
    images=os.listdir(new_dir)
    sum_test+=len(images)
print("测试集:{:}".format(sum_test))
        

训练集:

#统计训练集数目
import os
path="D:\\数据\\216_2021\\train"
sum_train=0
dirs=os.listdir(path)
for Dir in dirs:
    new_dir=os.path.join(path,Dir)
    images=os.listdir(new_dir)
    sum_train+=len(images)
print("训练集:{:}".format(sum_train))

总:

#统计总数目数目
import os
path="D:\\数据\\216_2021\\images"
sum=0
dirs=os.listdir(path)
for Dir in dirs:
    new_dir=os.path.join(path,Dir)
    images=os.listdir(new_dir)
    sum+=len(images)
print("总数:{:}".format(sum))

还好，最后训练集比测试集多一点。