数据集随机划分为训练集验证集测试集+数据集去后缀保存txt格式标签

时间：20220408
内容：数据集预处理

要把准备好的图片放到神经网络里训练，需要划分训练集，验证集，测试集。
今天找了几份代码不太好使，自己写了写。

首先是随机划分训练集，按照比例，比如8:1，那每轮9张图片。
每隔一轮生成一个随机数A，就取这一轮第A个为测试集。
随机划分为8:1的数据集：
代码：

import os
import shutil
from random import randrange
import os.path as osp

import mmcv


def divide_img(images_path):
    # 新建两个空文件夹
    save_train_path = 'labels/train'
    sava_test_path = 'labels/test'
    modnum = randrange(scale)  # 首先随机生成一个随机数
    # 开始划分训练集和测试集并保存
    # 遍历根文件夹
    for temp in os.listdir(images_path):  # temp是root的子文件夹
        image_dir = images_path + '/' + temp
        print(image_dir)  # 可以输出查看下该分类下的数据集是什么
        # 遍历子文件夹
        i = 0
        for imagename in os.listdir(image_dir):  # 获取具体的文件名称
            # 图片路径
            src_image = image_dir + '/' + imagename
            # print(modnum)
            if i % scale == modnum:  # 划分为测试集
                save_path = sava_test_path + '/'
                if not os.path.exists(save_path):
                    os.makedirs(save_path)
                shutil.copy(src=src_image, dst=save_path)
            else:  # 划分为训练集
                save_path = save_train_path + '/'
                if not os.path.exists(save_path):
                    os.makedirs(save_path)
                shutil.copy(src=src_image, dst=save_path)
                # 一轮结束
            if i % scale == 0 and i != 0:  # 将数据集划分为7：1，每次处理完八张图片后变化一次随机数
                modnum = randrange(scale)

            i = i + 1


if __name__ == '__main__':
    img_root1 = 'img'  # 图片的上上一层
    ann_dir = 'img/SegmentationClassPNG'  # 即图片所在文件夹 img/SegmentationClassPNG
    txt_path1 = 'divide'  # 后为img/divide
    # 随机生成训练集和测试集，比例为8：1
    scale = 9
    divide_img(img_root1)
    print("finish")

然后，训练还需要有文件名对应的标签，需要去掉后缀：

import mmcv
import os.path as osp

# 原始数据集 ann上一级
data_root = 'img'
# ann图像文件夹
ann_dir = "labels/test"
# txt文件保存路径
split_dir = 'img/divide'
mmcv.mkdir_or_exist(osp.join(data_root, split_dir))

filename_list = [osp.splitext(filename)[0] for filename in mmcv.scandir(
    osp.join(ann_dir), suffix='.png')]

with open(osp.join(split_dir, 'test.txt'), 'w') as f:
    # select first 4/5 as train set
    train_length = int(len(filename_list))
    f.writelines(line + '\n' for line in filename_list[:train_length])

如果小白不理解里面的文件夹格式，那么可以参考我的文件放置模式，然后改成自己的。

没有展开的文件夹里面都是图片。