【自动划分数据集】提供图片与标签（.txt）文件即可自动生成YOLO格式数据集

"""
input：
1.img_path 数据集图片地址
2.label_path 数据集标签地址（.txt格式）
3.划分比例（0到1） split_rate

output：
--dataset
    --images
        --train
        --val
    --labels
       --train
       --val
"""
import os
import random
import shutil

def get_imlist(path):
    return [os.path.join(path, f) for f in os.listdir(path) if f.endswith('.jpg')] #后缀要换成自己的

def move_labels(dest_dir, path):
    label_list = get_imlist(dest_dir + f'/images/{path}')
    for names in label_list:
        name = names.split('\')[-1:][0]
        name = name.split('.')[0]
        label = name + '.txt'
        if os.path.exists(f'{label_path}/'):
            shutil.copy(f'{label_path}/

', dest_dir + f'/labels/{path}') else: print(f'

is not exists') def main(src_path): dest_dir = output_path # 这个文件夹需要提前建好 img_list = get_imlist(src_path) random.shuffle(img_list) length = int(len(img_list) * split_rate) # 这个可以修改划分比例 os.makedirs(dest_dir + '/images/train') os.makedirs(dest_dir + '/images/test') os.makedirs(dest_dir + '/labels/train') os.makedirs(dest_dir + '/labels/test') for f in img_list[length:]: shutil.copy(f, dest_dir + '/images/train') for f in img_list[:length]: shutil.copy(f, dest_dir + '/images/test') #移动对应的标签到对应位置 move_labels(dest_dir, 'test') move_labels(dest_dir, 'train') print(f'finished') if __name__ == '__main__': path_dataset = r'..\output\crazing.15_0.5-0.5_1\expertLabel0.15_0.5' img_path = path_dataset + '\images' label_path = path_dataset + '\labels' split_rate = 0.2 output_path = r'datasets\demo' if not os.path.exists(output_path): os.makedirs(output_path) main(img_path)

https://54852.com/langs/942294.html