用python程序按比例划分数据集(训练集+测试集)

现有一个零售商品图片的数据集goods，没有做任何划分。goods目录下，子目录的名称为图片的类别，每个类别下有多张图片，文件结构如图1。现在，我们需要按一定的比例把它拆分为训练集和测试集。

图1. 数据集划分示意

实现该功能的思路是建立一个新的相同结构的目录树，并在原目录树中的每个类别中随机按比例地挑选部分图片，将它们移动到新目录树中相应的类别下。完整代码如下：

import os
import random
import shutil

#源数据集路径和目标数据集路径
path_source = './goods'
path_target = './goods_test'

#参数：源路径、目标路径和测试集所占比例
def seperate(path_source, path_target, percent):
    #生成包含path_source下所有目录名的列表
    categories = os.listdir(path_source)
    for name in categories:
        #在path_target下建立相同名称的子目录
        os.makedirs(os.path.join(path_target, name))
        #生成包含子目录下所有图片的列表
        nums = os.listdir(os.path.join(path_source, name))
        #随机按比例抽取一部分图片
        nums_target = random.sample(nums, int(len(nums)*percent))
        #把图片剪切到目标路径
        for pic in nums_target:
            shutil.move(os.path.join(path_source, name, pic), os.path.join(path_target, name, pic))

#执行完成后，path_source为训练集，path_target为测试集。
seperate(path_source, path_target, 0.3)

如果数据集是其他结构，在此基础上稍作改动即可。

其中的关键函数：

os.listdir(path)：列出path下所有项目的名称(包括目录和文件)

os.makedirs(path)：递归地创建path路径上的目录，目录已经存在的话会报错

os.makedir(path)可以用于仅创建一个目录的情况

random.sample(list, num)：从list中随机抽取num个元素组成新的列表

shutil.move(source, target)：将source文件移动到target文件

注意shutil.move是个很危险的操作，建议在调试代码时注释掉，确保代码无误后再执行。