拆分训练集,验证集

带json格式拆分训练集验证集

import os
import random
import shutil

def split_step_3():
    dir_a=r'F:\data\jupiter\det_no\wujian_det\train'

    img_files = ['%s/%s' % (i[0], j) for i in os.walk(dir_a) for j in i[-1] if j.endswith(('jpg', 'png', 'jpeg'))]

    random.shuffle(img_files)

    total = 8 + 2 + 5  # 总比例
    length = len(img_files)
    part1 = int(length * 8 / total)
    part2 = int(length * 2 / total)
    list2 = img_files[part1:part1 + part2]
    list3 = img_files[part1 + part2:]

    gropus=[list2,list3]
    print(len(list2),len(list3))

    targets=[r'F:\data\val',
             r'F:\data\test']
    copy_mode=1
    for index, group in enumerate(gropus):

        target_dir = targets[index]
        os.makedirs(target_dir, exist_ok=True)
        print(index, target_dir)
        for img_file in group:
            json_path = img_file.replace('.jpg', '.json')
            if os.path.exists(json_path):
                if copy_mode == 1:
                    shutil.copy(json_path, target_dir)
                if copy_mode == 2:
                    shutil.move(json_path, target_dir)
            if copy_mode == 1:
                shutil.copy(img_file, target_dir)
            if copy_mode == 2:
                shutil.move(img_file, target_dir)

def to_val():
    dir_a=r'F:\data\wujian_det'

    img_files = ['%s/%s' % (i[0], j) for i in os.walk(dir_a) for j in i[-1] if j.endswith(('jpg', 'png', 'jpeg'))]

    random.shuffle(img_files)

    length = len(img_files)
    part1 = int(length * 2 / 10)
    list1 = img_files[:part1]

    target_dir=r'F:\data\val'
    os.makedirs(target_dir, exist_ok=True)
    copy_mode=2
    for index, img_file in enumerate(list1):
        json_path = img_file.replace('.jpg', '.json')
        if os.path.exists(json_path):
            if copy_mode == 1:
                shutil.copy(json_path, target_dir)
            if copy_mode == 2:
                shutil.move(json_path, target_dir)
        if copy_mode == 1:
            shutil.copy(img_file, target_dir)
        if copy_mode == 2:
            shutil.move(img_file, target_dir)
if __name__ == '__main__':
    to_val()

你可能感兴趣的:(python基础,算法)