原始数据集划分训练集和测试集

原始数据集划分训练集和测试集

import os
import random
import shutil
import sys

#将所有原图与标签mask文件放在同一个文件夹里，运行该文件
#随后文件夹中通过不同文件类型来分离原图与mask文件

def movefile(filedir, tardir):
    pathdir = os.listdir(filedir)  # 取数据集的原始路径
    filenumber = len(pathdir)
    rate = 0.8  # 自定义抽取数据集的比例，比方说1000张抽100张，那就是0.1
    picknumber = int(filenumber * rate)  # 按照rate比例从数据集中取一定数量文件
    samples = random.sample(pathdir, picknumber)  # 随机选取picknumber数量的文件
    count = 0 #计数变量
    for sample in samples:# 当前目录下选取的数据
        if sample.endswith(".jpg"): # 查看选取的是否是jpg图片
            shutil.move(filedir + sample, tardir + sample)
            shutil.move(filedir + os.path.splitext(sample)[0] + '.png', tardir + os.path.splitext(sample)[0] + '.png')
            count += 1
        if(count / 0.1 > filenumber): # 移动的图片已经超过总数据集的10%，跳出
            break
    return


if __name__ == '__main__':
    filedir = r""  # 源数据集文件夹路径
    tardir = r''  # 移动到新的文件夹路径
    movefile(filedir, tardir)

# python random_select.py