YOLO类数据集划分

CangMing

YOLO类数据集划分代码

import os
import random
import shutil

# 定义文件路径
image_folder = "F:/Data_increase/images"
annotation_folder = "F:/Data_increase/yolo_txt"
train_folder = "F:/Data_increase/train"
train_xml_folder = "F:/Data_increase/train_txt"
val_folder = "F:/Data_increase/val"
val_xml_folder = "F:/Data_increase/val_txt"
test_folder = "F:/Data_increase/test"
test_xml_folder = "F:/Data_increase/test_txt"

# 创建保存文件夹
os.makedirs(train_folder, exist_ok=True)
os.makedirs(train_xml_folder, exist_ok=True)
os.makedirs(val_folder, exist_ok=True)
os.makedirs(val_xml_folder, exist_ok=True)
os.makedirs(test_folder, exist_ok=True)
os.makedirs(test_xml_folder, exist_ok=True)

# 获取所有图像文件的路径
image_files = [f for f in os.listdir(image_folder) if f.endswith('.jpg')]

# 随机打乱图像文件列表
random.shuffle(image_files)

# 计算划分数量
total_images = len(image_files)
train_split = int(0.7 * total_images)
val_split = int(0.2 * total_images)

# 划分并复制图像和对应的标注文件
for i, image_file in enumerate(image_files):
    annotation_file = image_file.replace('.jpg', '.txt')

    if i < train_split:
        shutil.copy(os.path.join(image_folder, image_file), os.path.join(train_folder, image_file))
        shutil.copy(os.path.join(annotation_folder, annotation_file), os.path.join(train_xml_folder, annotation_file))
    elif i < train_split + val_split:
        shutil.copy(os.path.join(image_folder, image_file), os.path.join(val_folder, image_file))
        shutil.copy(os.path.join(annotation_folder, annotation_file), os.path.join(val_xml_folder, annotation_file))
    else:
        shutil.copy(os.path.join(image_folder, image_file), os.path.join(test_folder, image_file))
        shutil.copy(os.path.join(annotation_folder, annotation_file), os.path.join(test_xml_folder, annotation_file))

dsax7

厉害 厉害

Cateners

楼主的代码整齐得像ChatGPT写的一样

CangMing

Cateners 发表于 2024-4-20 22:46
楼主的代码整齐得像ChatGPT写的一样

没有哦，这份代码是gpt出世之前写的