当我们已经标注好了目标检测的数据集（dota格式），并拥有相应的jpg格式的图像，我们可以通过代码构建一个用于模型训练的数据集。

代码说明

本脚本的目的是从已标注的图像和标签文件中创建一个组织好的目标检测数据集，用于后续的机器学习或者深度学习训练。数据集将被分为训练集（train）、验证集（validation）和测试集（test）。

输入

• 图像文件夹：包含所有.jpg格式图像的文件夹路径。
• 标签文件夹：包含所有对应图像标签的.txt文件（dota格式）的文件夹路径。
• 输出基础路径：存放分割后的数据集的根目录路径。

输出

• 在指定的输出基础路径下，将创建三个子文件夹：train、val 和 test，每个文件夹内部分别包含 images 和 labels 两个子文件夹。
  • train/images：训练集图像。
  • train/labels：训练集标签。
  • val/images：验证集图像。
  • val/labels：验证集标签。
  • test/images：测试集图像。
  • test/labels：测试集标签。

代码

"""
This code is used to create the data for training
"""

import os
import shutil
import random

def count_dotas_in_file(file_path):
    with open(file_path, 'r') as f:
        lines = f.readlines()
    return len(lines)  # 返回文件中目标的总数

def split_dataset(image_folder, label_folder, output_base):
    all_images = [img for img in os.listdir(image_folder) if img.endswith('.jpg')]
    random.shuffle(all_images)  # 随机排序图像

    # 初始化集合
    train_images, val_images, test_images = [], [], []
    total_dota_count = 0

    # 统计总的目标数量
    for img in all_images:
        label_file = img.replace('.jpg', '.txt')
        total_dota_count += count_dotas_in_file(os.path.join(label_folder, label_file))

    dota_count_so_far = 0  # 已处理好的目标数据

    # 分配图像到验证集和测试集
    for img in all_images:
        label_file = img.replace('.jpg', '.txt')
        current_file_dota_count = count_dotas_in_file(os.path.join(label_folder, label_file))
        # 按照6：2：2构建数据集
        if dota_count_so_far + current_file_dota_count <= 0.20 * total_dota_count:
            val_images.append(img)
        elif dota_count_so_far + current_file_dota_count <= 0.40 * total_dota_count:
            test_images.append(img)
        else:
            train_images.append(img)
        
        dota_count_so_far += current_file_dota_count

    # 创建输出文件夹并复制文件
    for subset, subset_images in zip(['train', 'val', 'test'], [train_images, val_images, test_images]):
        os.makedirs(os.path.join(output_base, subset, 'images'), exist_ok=True)
        os.makedirs(os.path.join(output_base, subset, 'labels'), exist_ok=True)
        for img in subset_images:
            shutil.copy(os.path.join(image_folder, img), os.path.join(output_base, subset, 'images', img))
            shutil.copy(os.path.join(label_folder, img.replace('.jpg', '.txt')),
                        os.path.join(output_base, subset, 'labels', img.replace('.jpg', '.txt')))


if __name__ == '__main__':
    image_folder = r'E:\project\代码\data_jpg'  # 图像文件夹路径
    label_folder = r'E:\project\代码\data_xml_jpg_v3'  # 标签文件夹路径
    output_base = r'E:\project\代码\train_v4\data_v1' # 输出文件夹路径
    split_dataset(image_folder, label_folder, output_base)

通过这个脚本，能够有效地准备好用于训练、验证和测试的数据集，以支持更复杂的机器学习应用。希望对大家有帮助！