YOLO测试集划分(多进程版)

介绍

GitHub:https://github.com/Eagle104fred/YOLOSplitDataset
本项目只用于分割labelimg生成的yolo格式的数据集, 默认进程池为逻辑核数量, 每个进程处理的任务个数为(单个数据集数量/逻辑核数量)
目录结构:

• root
  • A Dataset
    • images
    • labels
  • B Dataset
    • images
    • labels
  • C Dataset
    • images
    • labels

Getstarted

开始前请按照目录结构建立好文件夹和规范命名, 本程序放在root目录下运行, 生成的混合数据生成于 root/outFile 。

• 分割比率为 train：val => 9:1，需要修改的可以修改splitRatio这个参数。
• 如果不需要生成混合的train和val数据集请把（def CreateFilePath(self,doNotMixFlag=False): 下的doNotMixFlag设置为True，那样的话会分别为每个数据集生产image和label）

Warning: 程序运行前默认清除root/outFile, 请注意酌情备份!!!

Warning: 每个数据集下的文件名禁止同名, 否则会相互覆盖!!!

import os
import sys
import random
import shutil
from multiprocessing import Pool,cpu_count
import time

class RootFile:
    def __init__(self,path="./"):
        self.path=path

    def SplitSubRoot(self):
        def GetSubRootList(path):
            originDirList = os.listdir(path)

            originDirList.sort()
            filetedDirList = []
            for dir in originDirList:
                if dir == "outFile":
                    shutil.rmtree(os.path.join(path,dir))
                if os.path.isdir(os.path.join(path, dir)) == True and dir[0]!=".":
                    '''筛选出文件夹'''
                    filetedDirList.append(dir)
            return filetedDirList


        subRootList=GetSubRootList(self.path)

        for subRoot in subRootList:
            print("=========================================")
            print(subRoot + " Spliting")
            print("=========================================")

            subSplit = CreateValFile(subRoot)
            subSplit.Split()



class CreateValFile:
    def __init__(self,path="./",saveFile="outFile",splitRatio=0.1):
        self.path