【自存代码】划分数据集为训练集和测试集

文件夹我是手动建的，数据集名称目录下testA、testB，trainA、trainB，这四个文件夹下面再根据一共有多少类建子文件夹，从00开始编号，比如scribble是10类物体，就是00，01——10，qmul-sketch是3类，就是00，01，02。修改自这位大佬的代码：https://blog.csdn.net/weixin_47414034/article/details/12547

百里香酚兰

1897人浏览 · 2022-10-06 21:50:05

百里香酚兰 · 2022-10-06 21:50:05 发布

【自存代码】划分image和label为训练集和测试集

修改自这位大佬的代码：https://blog.csdn.net/weixin_47414034/article/details/125479366
因为自己需要划分scribble数据集，所以拿来魔改了一下
这是数据集拿到的样子：
在这里插入图片描述

这是魔改之后需要划分的样子：

在这里插入图片描述

文件夹我是手动建的，数据集名称目录下testA、testB，trainA、trainB，这四个文件夹下面再根据一共有多少类建子文件夹，从00开始编号，比如scribble是10类物体，就是00，01——10，qmul-sketch是3类，就是00，01，02
在这里插入图片描述

文件夹手动建完以后运行下面的代码就好了，可能我方法比较笨，但作为初学者的我来说这是唯一的办法
因为有10类数据要分，所以每次要把最后一级目录手动改成00，01，02，重复运行这段代码10次就行了
after_train_image=“scribble_dataset/scribble_dataset/scribble_dataset/trainA/10”
按道理来说这种重复的工作应该可以一步到位解决，但我目前还不会。

import random
import os
import shutil
#用八个变量来定义图片的存放路径
 
#划分前：
before_image = "scribble_dataset/scribble_dataset/scribble_dataset/images/basketball"
before_label = "scribble_dataset/scribble_dataset/scribble_dataset/scribbles/basketball"
 
#划分后的训练集
after_train_image="scribble_dataset/scribble_dataset/scribble_dataset/trainA/10"
after_train_label="scribble_dataset/scribble_dataset/scribble_dataset/trainB/10"

#划分后的测试集
after_test_image="scribble_dataset/scribble_dataset/scribble_dataset/testA/10"
after_test_label="scribble_dataset/scribble_dataset/scribble_dataset/testB/10"
 
# 所有的原始image的图片名字读入到一个列表里面
image= os.listdir(before_image)
#比如列表image=['1.jpg', '2.jpg', '3.jpg'.....]
 
#同理把原始label的图片名字读入到一个列表里面
label= os.listdir(before_label)
#比如列表label=['1.jpg', '2.jpg', '3.jpg'.....]
 
#读出image列表的长度：
length= len(image)
#生成一个[0,length-1]的列表，称之为索引列表
#比如列表index_list=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
index_list = list(range(length))
#打乱index_list列表
random.shuffle(index_list)
#此时index_list=[4, 0, 2, 7, 5, 8, 1, 6, 3, 9]
 
count=0
#把打乱过后的索引列表[4, 0, 2, 7, 5, 8, 1, 6, 3, 9]中每个索引拿出来：
for i in index_list:
    #比如说第一轮：image[i]=image[4],也就是说把image列表里第五张图片”5.jpg”这个字符串取出来
    #文件夹和文件名进行拼接，得到文件的完整路径
    imageName = os.path.join(before_image, image[i])
    #比如imageName=“E:/桌面/before/image/5.jpg”
 
    #label图片也是同理，得到相应的label图片的完整路径
    labelName = os.path.join(before_label, label[i])
 
    #下面开始划分（注意这里采取的是复制图片然后再转移路径，而不是直接转移路径
    #也就是说划分完，原图文件夹里面还是有图片
    #这里采取：train：val：test=7：2：1的划分比例
    if count <length * 0.7:
    #copy函数：前面第一个参数是原始文件路径，第二个参数是复制到哪个文件夹里面
        shutil.copy2(imageName, after_train_image)
        shutil.copy2(labelName, after_train_label)
    else:
        shutil.copy2(imageName, after_test_image)
        shutil.copy2(labelName, after_test_label)
    count=count+1