现有一个零售商品图片的数据集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是个很危险的操作,建议在调试代码时注释掉,确保代码无误后再执行。

Logo

为开发者提供学习成长、分享交流、生态实践、资源工具等服务,帮助开发者快速成长。

更多推荐