目标检测——标签类别映射批量修改

对于两个数据集，第一个数据集的people是1号类别，第二个数据集的people是9号类别，此时需要通过映射的方式去批量修改数据集，在对其进行合并。

示例：

我有两个数据集如下：
 

第一个数据集类别：

0: bacterial-blight
1: bacterial-leaf-streak
2: blast
3: brown-spot
4: dead-heart
5: downy-mildew
6: false-smut
7: normal
8: sheath-blight
9: tungro

第二个数据集类别：

0: Bacteria Leaf Blight
1: Brown Spot
2: Leaf smut

• 第二个数据集的 Bacteria Leaf Blight（类别 0）对应第一个数据集的 bacterial-blight（类别 0）。

• 第二个数据集的 Brown Spot（类别 1）对应第一个数据集的 brown-spot（类别 3）。

• 第二个数据集的 Leaf smut（类别 2）对应第一个数据集的 false-smut（类别 6）

代码：
 

import os

# 建立类别映射关系
category_mapping = {
    0: 0,  # Bacteria Leaf Blight -> bacterial-blight
    1: 3,  # Brown Spot -> brown-spot
    2: 6   # Leaf smut -> false-smut
}

# 输入和输出文件夹路径
input_folder = "path/to/second_dataset_labels"  # 第二个数据集的标签文件夹路径
output_folder = "path/to/modified_labels"       # 修改后的标签文件夹路径

# 确保输出文件夹存在
os.makedirs(output_folder, exist_ok=True)

# 遍历输入文件夹中的所有标签文件
for filename in os.listdir(input_folder):
    if filename.endswith(".txt"):
        input_path = os.path.join(input_folder, filename)
        output_path = os.path.join(output_folder, filename)
        
        # 读取标签文件
        with open(input_path, "r") as f:
            lines = f.readlines()
        
        # 修改类别号
        modified_lines = []
        for line in lines:
            parts = line.strip().split()
            if len(parts) >= 1:
                # 获取类别号并映射到新的类别号
                original_category = int(parts[0])
                if original_category in category_mapping:
                    new_category = category_mapping[original_category]
                    modified_line = f"{new_category} " + " ".join(parts[1:])
                    modified_lines.append(modified_line)
                else:
                    # 如果类别号不在映射中，保持原样
                    modified_lines.append(line.strip())
        
        # 写入修改后的标签文件
        with open(output_path, "w") as f:
            f.write("\n".join(modified_lines))

print("类别号修改完成！")