yolo格式数据集划分代码

yolo格式数据集划分代码

一、描述

目前只有两个文件夹，image—存放图像，labels----存放txt格式的标签文件。
需要划分为训练集、验证集和测试集。
手动划分的话，一个费时费力，还无法保证随机性。

二、按图像划分代码

####从图像文件夹中随机抽取一定比例的图像，并根据图像名字，抽取对应的txt标签
import os
import random
import shutil
def moveFile(Imgdir):
        pathDir = os.listdir(Imgdir)    #获取图像列表
        filenumber=len(pathDir)
        rate=0.1    #定义抽取图片的比例
        random_imgnum=int(filenumber*rate) #按照比例从文件夹中取一定数量图片
        sample = random.sample(pathDir, random_imgnum)  #随机选取random_imgnum数量的样本图片
        print (sample)
        for imgname in sample:
                #name0 = os.path.join(outdir1,os.path.basename(name))
                nametxt=os.path.splitext(imgname)[0]
                shutil.move(Imgdir+imgname, img_target_dir+imgname)
                shutil.move(Txtdir+nametxt+'.txt', Txt_target_dir+nametxt+'.txt')
        return


if __name__ == '__main__':
    
    Imgdir = "/home/jun/work/img/"#图像文件夹
    img_target_dir = "/home/jun/work/img1/"#划分目标图像文件夹
    Txtdir = "/home/jun/work//txt/"#标签文件夹
    Txt_target_dir="/home/jun/work/txt1/"#划分目标标签文件夹
    moveFile(Imgdir)

三、按标签划分代码

####从txt标签文件夹中随机抽取一定比例的txt标签，并根据txt标签名字，抽取对应的图像
import os
import random
import shutil
def moveFile(Txtdir):
        pathDir = os.listdir(Txtdir)    #获取标签列表
        filenumber=len(pathDir)
        rate=0.2    #定义抽取图片的比例
        random_imgnum=int(filenumber*rate) #按照比例从文件夹中取一定数量标签
        sample = random.sample(pathDir, random_imgnum)  #随机选取random_imgnum数量的样本标签
        print (sample)
        for labelname in sample:
            nameimg = os.path.splitext(labelname)[0]
            shutil.move(Txtdir+labelname, Txt_target_dir+labelname)
            shutil.move(Imgdir+nameimg+'.jpg', img_target_dir+nameimg+'.jpg')

        '''
        for imgname in sample:
                #name0 = os.path.join(outdir1,os.path.basename(name))
                nametxt=os.path.splitext(imgname)[0]
                shutil.move(Imgdir+imgname, img_target_dir+imgname)
                shutil.move(Txtdir+nametxt+'.txt', Txt_target_dir+nametxt+'.txt')
        '''
        return


if __name__ == '__main__':
    
    Imgdir = "/images/"#图像文件夹
    img_target_dir = "/img/"#划分目标图像文件夹
    Txtdir = "/labels/"#标签文件夹
    Txt_target_dir="/ann/"#划分目标标签文件夹
    #moveFile(Imgdir)
    moveFile(Txtdir)