将两个YOLO格式的数据集合并，并保持相同类别

1. 需求分析

最近在做两个YOLO格式的数据集合并，第一个数据集包含了第二个数据集的类别，但是相应的类别id对应不住，需要修改第二个数据集的类别标签与第一个数据集对应住。

2. 修改第二个数据集标签对应

2.1 实现思路

1. 导入所需的库：

   • os：用于操作文件和文件夹路径。
   • tqdm：用于在循环中显示进度条。

2. 定义类别 ID 映射字典 class_id_mapping：

   • 这个字典用于将原始类别 ID 映射到新的类别 ID。

3. 定义函数 convert_yolo_annotations(yolo_annotations_root, new_yolo_annotations_root)：

   • 这个函数接受两个参数，即原始 YOLO 格式的标注文件夹路径 yolo_annotations_root 和新的 YOLO 格式的标注文件夹路径 new_yolo_annotations_root。
   • 函数首先遍历 train、val、test 数据集分割文件夹。
   • 对于每个数据集分割，如果该分割不存在标注文件夹，则打印警告信息并跳过处理。
   • 如果标注文件夹存在，则遍历其中的每个文件。
   • 对于每个文件，打开原始 YOLO 格式的标注文件和新的 YOLO 格式的标注文件，并逐行读取原始标注文件。
   • 对于每一行，将其按空格分割为列表，获取原始类别 ID，并根据 class_id_mapping 字典将其替换为新的类别 ID。
   • 将修改后的行写入新的标注文件。
   • 在处理完成后，打印消息表示类别 ID 映射完成，并保存为新的 YOLO 格式的标注文件。

4. 调用 convert_yolo_annotations() 函数：

   • 最后，调用 convert_yolo_annotations() 函数，传入了原始 YOLO 格式的标注文件夹路径和新的 YOLO 格式的标注文件夹路径。

2.2 实现代码

import os
from tqdm import tqdm

# 创建类别 ID 映射字典
class_id_mapping = {
    0: 0,  # 将原始类别 ID 为 0 的类别映射为新的类别 ID 为 0
    1: 13,  # 将原始类别 ID 为 1 的类别映射为新的类别 ID 为 13
    # 添加更多的映射关系...
}

# 定义原始 YOLO 格式的标注文件夹路径和新的 YOLO 格式的标注文件夹路径
yolo_annotations_root = r'D:\data\AIR-MOT\labels'
new_yolo_annotations_root = r'D:\data\AIR-MOT\labels_new'

def convert_yolo_annotations(yolo_annotations_root, new_yolo_annotations_root):
    # 遍历 train、val、test 数据集分割文件夹
    for dataset_type in ['train', 'val', 'test']:
        # 拼接原始 YOLO 格式的标注文件夹路径
        yolo_annotations_folder = os.path.join(yolo_annotations_root, dataset_type)

        # 如果当前数据集分割中没有标注文件夹，则跳过处理
        if not os.path.exists(yolo_annotations_folder):
            print(f"警告：{dataset_type} 数据集分割中没有标注文件夹，跳过处理。")
            continue

        # 拼接新的 YOLO 格式的标注文件夹路径
        new_yolo_annotations_folder = os.path.join(new_yolo_annotations_root, dataset_type)

        # 确保新的 YOLO 格式的标注文件夹存在
        if not os.path.exists(new_yolo_annotations_folder):
            os.makedirs(new_yolo_annotations_folder)

        # 遍历原始 YOLO 格式的标注文件夹中的所有文件
        for filename in tqdm(os.listdir(yolo_annotations_folder), desc="新类别标签转换"):
            # 拼接原始 YOLO 格式的标注文件路径
            original_file_path = os.path.join(yolo_annotations_folder, filename)

            # 拼接新的 YOLO 格式的标注文件路径
            new_file_path = os.path.join(new_yolo_annotations_folder, filename)

            # 打开原始 YOLO 格式的标注文件和新的 YOLO 格式的标注文件
            with open(original_file_path, 'r') as f, open(new_file_path, 'w') as f_new:
                # 逐行读取原始标注文件
                for line in f:
                    # 将每行按空格分割成列表
                    parts = line.strip().split()

                    # 获取原始类别 ID
                    class_id = int(parts[0])

                    # 如果原始类别 ID 在映射字典中，则替换为新的类别 ID
                    if class_id in class_id_mapping:
                        parts[0] = str(class_id_mapping[class_id])

                    # 将修改后的行写入新的标注文件
                    f_new.write(' '.join(parts) + '\n')

        print(f"类别 ID 映射完成，并保存为新的 YOLO 格式的标注文件，数据集类型：{dataset_type}。")

convert_yolo_annotations(yolo_annotations_root, new_yolo_annotations_root)

3. 合并两个数据集

将两个不同格式的数据集合并成一个符合YOLO格式的数据集。具体流程如下：

1. 创建一个新的目录作为合并后的数据集根目录。
2. 遍历数据集的不同类型（如训练集、验证集、测试集）。
3. 对于每种类型的数据集：
   • 创建符合YOLO格式的目录结构（包括images和labels子目录）。
   • 将第一个数据集中的图像文件和标签文件复制到合并后的数据集目录中。
   • 然后将第二个数据集中的图像文件和标签文件复制到合并后的数据集目录中，但会避免覆盖已存在的文件。
   • 生成符合YOLO格式的索引文件，将图像文件路径写入索引文件中。

这样一来，就能够将两个不同格式的数据集合并成一个符合YOLO格式的数据集，并且保证了不会覆盖已有的文件。
数据集格式：

import os
from tqdm import tqdm
import shutil


def copy_files(src_dir, dst_dir):
    """
    复制源目录中的所有文件到目标目录。
    """
    if not os.path.exists(dst_dir):
        os.makedirs(dst_dir)

    if not os.path.exists(src_dir):
        print(f"警告：源文件夹 {src_dir} 不存在，无法复制任何文件。")
        return

    # 检查源目录是否为空
    src_files = os.listdir(src_dir)
    if not src_files:
        print(f"警告：源文件夹 {src_dir} 是空的，无法复制任何文件。")
        return

    for filename in tqdm(os.listdir(src_dir), desc=f'Copying {src_dir}'):
        src_path = os.path.join(src_dir, filename)
        dst_path = os.path.join(dst_dir, filename)
        if not os.path.exists(dst_path):
            shutil.copy2(src_path, dst_path)
        elif os.path.isdir(dst_path):
            print(f"警告：目标路径 {dst_path} 已经存在且是一个文件夹，跳过复制文件 {filename}。")
        else:
            print(f"警告：目标文件 {dst_path} 已经存在，跳过复制文件 {filename}。")


def merge_yolo_datasets(dataset1_root, dataset2_root, merged_dataset_root):
    # 创建合并后的数据集根目录
    if not os.path.exists(merged_dataset_root):
        os.makedirs(merged_dataset_root)

    # 定义数据集类型列表
    dataset_types = ['train', 'val', 'test']

    # 遍历数据集类型
    for dataset_type in tqdm(dataset_types, desc="创建合并后数据集"):

        # 创建符合 YOLO 格式的目录结构
        os.makedirs(os.path.join(merged_dataset_root, 'images', dataset_type), exist_ok=True)
        os.makedirs(os.path.join(merged_dataset_root, 'labels', dataset_type), exist_ok=True)

        # 复制第一个数据集的图像文件和标注文件到合并后的数据集目录
        src_dataset1_images = os.path.join(dataset1_root, 'images', dataset_type)
        dst_dataset_images = os.path.join(merged_dataset_root, 'images', dataset_type)
        copy_files(src_dataset1_images, dst_dataset_images)

        src_dataset1_labels = os.path.join(dataset1_root, 'labels', dataset_type)
        dst_dataset_labels = os.path.join(merged_dataset_root, 'labels', dataset_type)
        copy_files(src_dataset1_labels, dst_dataset_labels)

        # 复制第二个数据集的图像文件和标注文件到合并后的数据集目录，避免覆盖第一个数据集的文件
        src_dataset2_images = os.path.join(dataset2_root, 'images', dataset_type)
        copy_files(src_dataset2_images, dst_dataset_images)

        src_dataset2_labels = os.path.join(dataset2_root, 'labels', dataset_type)
        copy_files(src_dataset2_labels, dst_dataset_labels)

        # 生成合并后的数据集类型的索引文件
        with open(os.path.join(merged_dataset_root, 'labels',  f'{dataset_type}.txt'), 'w') as f:
            for image_file in os.listdir(dst_dataset_images):
                f.write(f"{os.path.join('images', image_file)}\n")

    print("数据集合并完成！")

# 调用函数进行数据集合并
merge_yolo_datasets(r'D:\data\DIOR\yolo', r'D:\data\AIR-MOT', r'D:\data\mot')