[人工智能]2021-08-28 YOLO格式数据集划分

YOLO格式的数据集划分

标题标签格式txt

import os
import cv2
import random
from PIL import Image
# os.makedirs('labels_train')
# os.makedirs('labels_val')

img_train_path ='./img_train/'  #图片训练集保存路径
img_val_path = './img_val/'     #图片验证集保存路径
labels_train_path = './labels_train/'       #标签训练集保存路径
labels_val_path = './labels_val/'           #标签验证集保存路径
dir_images = './img/'                       #数据集路径
dir_labels  = './ann-txt/'                  #标签路径

images = os.listdir(dir_images)             #把所有的图片名放入一个列表中
labels = os.listdir(dir_labels)

random.seed(2021)                           #设置一个随机种子，确保每次运行都按照既定的随机形式
random.shuffle(images)                      #洗牌操作，打乱列表顺序
random.seed(2021)                           #再一次进行洗牌操作时要重新设置随机种子
random.shuffle(labels)
train_spilt_rate =0.8       #数据集划分比例
a = int(len(images)*train_spilt_rate)   #训练集数量

for i,image in enumerate(images):   #对图片进行划分
    image_path = dir_images + image
    img = Image.open(image_path)    #读取图片
    if i < a:
        img.save(img_train_path + image)    #保存训练集
    else:
        img.save(img_val_path + image)  #保存验证集

for i , label in enumerate(labels):     #对标签进行划分
    with open(dir_labels + label, 'r') as f:    #读取txt格式的标签
        if i < a:
            with open(labels_train_path + label , 'w') as s:    #保存训练集
                s.write(f.read())
        else:
            with open(labels_val_path + label , 'w') as s:      #保存验证集
                s.write(f.read())