划分训练集、验证集和测试集代码

输入需要划分的数据所在的文件夹，返回一个划分好的包含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]