voc数据集划分训练集与测试集（nanodet训练用）

voc数据集划分训练集与测试集（nanodet训练用）

import glob
import os
import random
import shutil

from PIL import Image


def split_dataset(img_path, ann_path, train_dir, val_dir, test_dir=None, is_split_test=False):
    imgs, masks = get_image(img_path, ann_path)  # get_image在上文中有

    # 创建一些目录
    train_imgs_dir = os.path.join(train_dir, 'imgs')
    train_anns_dir = os.path.join(train_dir, 'anns')
    val_imgs_dir = os.path.join(val_dir, 'imgs')
    val_anns_dir = os.path.join(val_dir, 'anns')
    mkdir(train_imgs_dir)
    mkdir(train_anns_dir)
    mkdir(val_imgs_dir)
    mkdir(val_anns_dir)
    if test_dir:
        test_imgs_dir = os.path.join(train_dir, 'imgs')
        test_anns_dir = os.path.join(train_dir, 'anns')
        mkdir(test_imgs_dir)
        mkdir(test_anns_dir)

    #  assert len(masks) == len(imgs), 'len(imgs_dir) and len(masks_dir) must be the same'
    indexs = [i for i in range(len(imgs))]
    print(len(imgs))
    random.shuffle(indexs)  # 打乱索引
    print(indexs)
    print(type(indexs))

    if is_split_test:  # 若是要划分测试集，则比例为6：2：2
        cnt_train = int(round(len(imgs) * 0.8, 0))  # 四舍五入，0表示保留0个小数点
        print(cnt_train)
        cnt_val = int(round(len(imgs) * 0.2, 0))  # 因为在需要做索引，所以强制转换为整型
        print(cnt_val)
        cnt_test = int(round(len(imgs) * 0.0, 0))

        # 下面使用列表切片得到train/val/test数据的索引
        train_indexs = indexs[0:cnt_train]
        print(train_indexs)
        val_indexs = indexs[cnt_train:(cnt_train + cnt_val):1]
        test_indexs = indexs[(cnt_train + cnt_val):]

    else:  # 只划分训练集与验证集, 这里使用了 7:3
        cnt_train = int(round(len(imgs) * 0.7, 0))  # 四舍五入，0表示保留0个小数点
        cnt_val = int(round(len(imgs) * 0.3, 0))  # 因为在需要做索引，所以强制转换为整型
        train_indexs = indexs[0:cnt_train]
        print(train_indexs)
        val_indexs = indexs[cnt_train:]

    # 保存图片
    for i in train_indexs:
        # imgs[i].save(os.path.join(train_imgs_dir, imgs[i]))
        # # masks[i].save(os.path.join(args.train_dir, 'annotation', mask_names[i]))
        # masks[i].save(os.path.join(train_anns_dir, imgs[i][:-4] + '.xml'))
        shutil.copy(imgs[i], train_imgs_dir)
        shutil.copy(os.path.join(masks[i][:-4] + '.xml'), train_anns_dir)
    for i in val_indexs:
        shutil.copy(imgs[i], val_imgs_dir)
        shutil.copy(os.path.join(masks[i][:-4] + '.xml'), val_anns_dir)
    if is_split_test:
        for i in test_indexs:
            imgs[i].save(os.path.join(test_imgs_dir, imgs[i]))
            # masks[i].save(os.path.join(args.val_dir, 'annotation', mask_names[i]))
            masks[i].save(os.path.join(test_anns_dir, masks[i][:-4] + '.xml'))


def mkdir(path):  # 用于创建目录
    # 引入模块
    # 去除首位空格
    path = path.strip()
    # 去除尾部 \ 符号
    path = path.rstrip("\\")
    # 判断路径是否存在
    # 存在     True
    # 不存在   False
    isExists = os.path.exists(path)
    # 判断结果
    if not isExists:
        # 如果不存在则创建目录
        # 创建目录操作函数
        os.makedirs(path)
        return True
    else:
        # 如果目录存在则不创建，并提示目录已存在
        return False


# 下面是我自定义的函数，用于只取出图片用测试，并且把对应的名字取出来用保存最后的结果。用于water数据集
# 下面是我自定义的函数，用于只取出图片用测试，并且把对应的名字取出来用保存最后的结果。用于water数据集
def get_image(img_path, ann_path):
    img_list = glob.glob(img_path + r'\*')
    mask_list = glob.glob(ann_path + r'\*')
    # print(len(img_list))
    # print('mask_list', len(mask_list))
    # imgs = []
    # names = []
    # masks = []
    #
    # for img_path in img_list:
    #     names.append(os.path.basename(img_path))
    #     img = Image.open(img_path)
    #     imgs.append(img)
    # for mask_path in mask_list:
    #     mask = Image.open(mask_path)  # 读取图片，读取
    #     masks.append(mask)
    return img_list, mask_list


if __name__ == "__main__":
#这里是自己的voc数据集的位置，前面是全部图片的位置，后面是全部XML文件的位置
    split_dataset("G:\\laboratory\\Experimentrelated\\Estrus\\Postures\\datasets\\all image","G:\\laboratory\\Experimentrelated\\Estrus\\Postures\\datasets\\all xml","train","val")