博客

利用txt文件划分训练集、测试集与验证集

在基于深度学习的目标检测任务中，往往需要将数据图片划分成训练集、验证集与测试集，在此记录一种笨办法，各位大佬可忽略。

训练集：日常学习

验证集：模拟考试

测试集：期末考试

• 第一步，先将原jpg数据图片与标签txt文件分别存放在一个文件夹内，如下图：
  
  
  
  然后将以下代码编辑在一个py文件内，eg:data_class.py:

# 将一个文件夹下数据按比例分在三个文件夹下
import os
import random
import shutil
from shutil import copy2
datadir_normal = "/****/*****/****/****/JPEGImages/labels"#更换为自己的路径

all_data = os.listdir(datadir_normal)#（图片文件夹）
num_all_data = len(all_data)
print( "总数量: " + str(num_all_data) )
index_list = list(range(num_all_data))
print(index_list)

#生成随机种子，使得每一次划分的数据一样，如果没有要求，可以删除本段代码
r=random.random
random.seed(2021)
random.shuffle(index_list,random=r)
print(index_list)
#生成随机种子，使得每一次划分的数据一样，如果没有要求，可以删除本段代码

num = 0
trainDir = "/****/*****/****/****/JPEGImages/train/"#（将训练集放在这个文件夹下）
if not os.path.exists(trainDir):
    os.mkdir(trainDir)     
validDir = '/****/*****/****/****/JPEGImages/valid/'#（将验证集放在这个文件夹下）
if not os.path.exists(validDir):
    os.mkdir(validDir)        
testDir = '/****/*****/****/****/JPEGImages/test/'#（将测试集放在这个文件夹下）        
if not os.path.exists(testDir):
    os.mkdir(testDir)
        
for i in index_list:
    fileName = os.path.join(datadir_normal, all_data[i])
    if num < num_all_data*0.6: #Train:0.6
        print(str(fileName))
        copy2(fileName, trainDir)
    elif num>num_all_data*0.6 and num < num_all_data*0.8:#valid:0.2
        print(str(fileName))
        copy2(fileName, validDir)
    else:#test:0.2
        copy2(fileName, testDir)
    num += 1

运行后，将在JPEGImages文件夹下看到多了train、test、valid3个文件夹，如下图：



文件夹内容就是根据上边的程序划分的标签txt文件。

• 第二步，提取刚刚生成的train、test、valid3个文件夹内的txt文件名称并分别存放在TRAIN、TEST、VALID3个txt文件中。
• 提取命令：
• 
find /*****/******/*****/****/JPEGImages/test -name "*.txt" > TEST.txt

• 
find /*****/******/*****/****/JPEGImages/train -name "*.txt" > TRAIN.txt

• 
find /*****/******/*****/****/JPEGImages/valid -name "*.txt" > VALID.txt

  
  三条命令运行结束以后可以看到多了三个txt文件，如下图：
  
  
  
  打开文件，内容为：

/home/outsider/其他/1/JPEGImages/valid/20210110113149061.txt

需要使用文档的查找与替换功能将3个txt文档内的名称修改为如下格式：

20210110113149061

文档的查找与修改如下：



替换后效果如下：

• 第三步，将以下代码编辑在一个py文件内，eg:copy_class.py:

from PIL import Image
im_num = []
for line in open("/***/****/***/VALID.txt", "r"):
    im_num.append(line)
    print(im_num)

for a in im_num:
    im_name = '/***/***/***/JPEGImages/images/{}'.format(a[:-1]) + '.jpg'  #原始路径
    print(im_name)
    im = Image.open(im_name)  # 打开指定路径下的图像

    tar_name = '/***/***/***/JPEGImages/valid/{}'.format(a[:-1]) + '.jpg'  #移动后的路径
    print(tar_name)
    im=im.convert('RGB')
    im.save(tar_name)  # 另存
    im.close()

每次运行后，需要将文件名修改下，修改为相对应的txt文档，程序运行结束后，将在train、test、valid3个文件夹内看到程序已经将标签txt对应的图片复制过来了。

到此，数据集就已经划分好了，也可以根据图片名称将对应的txt文件复制到相应的文件夹内，方法比较老实，大佬们就不用看啦。

最后，我将两个py文件放在下方。