输入需要划分的数据所在的文件夹,返回一个划分好的包含train,val和test的文件夹。
val_ratio,和test_ratio分别为验证集和测试集所占的比例,test_ratio=0表示不划分测试集。
import os
import cv2
import random
import shutil
import numpy as np
from tqdm import tqdm
from collections import defaultdict
'''
源文件夹下的图片和标签文件需要命名为images和labels
记得修改image的扩展名:jpg or png
'''
def main(src_dir, dst_dir, val_ratio=0.1, test_ratio=0):
src_imgs_dir = os.path.join(src_dir, 'images')
src_labels_dir = os.path.join(src_dir, 'labels')
fids = [f for f in os.listdir(src_imgs_dir) if f.endswith('.png')]
random.seed(6)
random.shuffle(fids)
total_num = len(fids)
val_num = int(total_num * val_ratio)
test_num = int(total_num * test_ratio)
train_num = total_num - val_num - test_num
train_set = fids[:train_num]
val_set = fids[train_num: train_num + val_num]