前言
深度学习说到底是由数据驱动的，所以数据是非常重要的。我们在网上收集的数据，常常是没有分成训练集，验证集和测试的，需要我们自己进行分割。本笔记的代码参照了深度之眼老师提供的参考代码，并作了一些相关的拓展。
代码和数据集：dataSplit.zip
任务
在网上收集不同类别的数据，笔者收集了两个类别的数据，ants和bees各100张，放在old_data下的两个文件夹内。拆分数据集为训练集，验证集和测试集。
数据集拆分
拆分数据集，即要随机分配数据到训练集，验证集和测试集。可以有两种方式，第一种是直接拆分数据，即直接在物理地址上切分；第二种是分别生成训练集、验证集和测试集的路径列表文件，然后在读取过程中切分。
- 要用到的函数简析
os.listdir()： 获取给定路径下的文件和文件夹列表，不能得到下一级目录的文件和文件夹列表。如：
print(os.listdir(old_path))
# output
# ['ants', 'bees']
os.walk()： 该方法有4个参数，我们就用一个就可以，即top，是要遍历的路径。该方法可以遍历完给定路径下的所有文件信息，为定路径下的每一个目录（包括自身）返回的一个三元组，即(root, dirs , files)。root指当前遍历文件夹的的地址，dirs 是一个 list ，内容是该文件夹中所有的目录的名字(不包括子目录)，files 同样是 list , 内容是该文件夹中所有的文件。如：
root_dir, sub_dirs, sub_sub_dirs = os.walk(old_path)
print(root_dir, '\n', sub_dirs, '\n', sub_dirs1)
# output
# ('old_data', ['ants', 'bees'], [])
# ('old_data\\ants', [], ['1030023514_aad5c608f9.jpg', '1095476100_3906d8afde.jpg', .....])
# ('old_data\\bees', [], ['1092977343_cb42b38d62.jpg', '1093831624_fb5fbe2308.jpg', .....]
for root_dir, sub_dirs, files in os.walk(old_path): # 遍历os.walk(）返回的每一个三元组，内容分别放在三个变量中
print('root_dir:', root_dir, 'sub_dirs:', sub_dirs, 'file:', files)
# output
# root_dir: old_data sub_dirs: ['ants', 'bees'] file: []
# root_dir: old_data\ants sub_dirs: [] file: ['1030023514_aad5c608f9.jpg', .....]
# root_dir: old_data\bees sub_dirs: [] file: ['1092977343_cb42b38d62.jpg', .....]
os.path.join()： 把字符拼接成目录。
shuffle() ： 将序列的所有元素随机排序
filter():filter(function, iterable)，用于过滤序列，过滤掉不符合条件的元素，返回一个迭代器对象，如果要转换为列表，可以使用 list() 来转换。该函数接收两个参数，第一个为判断函数，第二个为序列，序列的每个元素作为参数传递给函数进行判断，然后返回 True 或 False，最后将返回 True 的元素放到新列表中。如下表示去掉img_names列表中非.jpg文件，并将结果转换成列表。
list(filter(lambda x: x.endswith('.jpg'), img_names))
- 第一种：直接切分数据
要用到shutil包来复制数据。思路就是遍历目录下的所有文件，得到所有类别的存放路径。遍历每个类别，然后获得每个类别的所有图片的名称列表，并将其随机排列。然后创建存放对应数据集的文件夹，并遍历图片名称列表，按照比例（8:1:1）将名字分配给不同的集。最后拼接路径，根据路径复制图片到对应位置。代码如下：
import os
import random
import math
import shutil
def data_split(old_path):
new_path = 'data'
if os.path.exists('data') == 0:
os.makedirs(new_path)
for root_dir, sub_dirs, file in os.walk(old_path): # 遍历os.walk(）返回的每一个三元组，内容分别放在三个变量中
for sub_dir in sub_dirs:
file_names = os.listdir(os.path.join(root_dir, sub_dir)) # 遍历每个次级目录
file_names = list(filter(lambda x: x.endswith('.jpg'), file_names)) # 去掉列表中的非jpg格式的文件
random.shuffle(file_names)
for i in range(len(file_names)):
if i < math.floor(0.8*len(file_names)):
sub_path = os.path.join(new_path, 'train_set', sub_dir)
elif i < math.floor(0.9*len(file_names)):
sub_path = os.path.join(new_path, 'val_set', sub_dir)
elif i < len(file_names):
sub_path = os.path.join(new_path, 'test_set', sub_dir)
if os.path.exists(sub_path) == 0:
os.makedirs(sub_path)
shutil.copy(os.path.join(root_dir, sub_dir, file_names[i]), os.path.join(sub_path, file_names[i])) # 复制图片，从源到目的地
if __name__ == '__main__':
data_path = 'old_data'
data_split(data_path)
结果：
第二种：生成路径列表
代码与上述方法类似，此方法只是将路径，分别储存到三个不同的文本文件中。这种方法可以直接生成路径列表，在加载数据时不用重新生成，使得数据读取过程简化了。代码如下：
def data_list(path):
for root_dir, sub_dirs, _ in os.walk(path): # 遍历os.walk(）返回的每一个三元组，内容分别放在三个变量中
idx = 0
for sub_dir in sub_dirs:
file_names = os.listdir(os.path.join(root_dir, sub_dir)) # 遍历每个次级目录
file_names = list(filter(lambda x: x.endswith('.jpg'), file_names)) # 去掉列表中的非jpg格式的文件
random.shuffle(file_names)
for i in range(len(file_names)):
if i < math.floor(0.8 * len(file_names)):
txt_name = 'train_set.txt'
elif i < math.floor(0.9 * len(file_names)):
txt_name = 'val_set.txt'
elif i < len(file_names):
txt_name = 'test_set.txt'
with open(os.path.join(path, txt_name), mode='a') as file:
file.write(str(idx) + ',' + os.path.join(path, sub_dir, file_names[i]) + '\n') # 为了以后好用，修改了这里，将' '改成了','，另外路径加了sub_dir
idx += 1
if __name__ == '__main__':
data_path = 'old_data'
data_list(data_path)
结果：
参考
https://ai.deepshare.net/detail/p_5df0ad9a09d37_qYqVmt85/6
https://www.runoob.com/python/os-walk.html
https://www.runoob.com/python/os-listdir.html