深度学习中的自动化标签转换：对数据集所有标签做映射转换

news2024/9/24 19:25:11

在这里插入图片描述
在机器学习中，特别是在涉及图像识别或分类的项目中，标签数据的组织和准确性至关重要。本文探讨了一个旨在高效转换标签数据的 Python 脚本。该脚本在需要更新或更改类标签的场景中特别有用，这是正在进行的机器学习项目中的常见任务。我们将逐步介绍如何使用此脚本并了解其功能。

使用脚本的分步指南

初始设置：
脚本首先导入必要的模块：**os用于文件操作、shutil用于高级文件操作、zipfile用于处理 zip 文件以及datetime**用于时间戳。

定义路径：
您需要指定标签目录的路径。然后，该脚本会自动在同一目录中创建备份 zip 文件路径，并带有时间戳以确保唯一性。

创建备份：
在进行任何更改之前，该脚本会将现有标签文件备份到 zip 存档中。这就像在进行任何更改之前拍摄数据快照一样，确保您有后备选项。

标签转换：
核心功能涉及更新标签文件中的类索引。您定义当前和新的类名称，脚本将它们映射到它们各自的索引。此过程类似于分类系统更改时更新图书馆中的目录。

应用更改：
脚本迭代每个标签文件，应用映射来更新类索引。这就像检查文件柜并更新每个文件夹上的标签。

输出和验证：
提供更改的摘要，包括修改的文件数量以及修改前后最后一个文件内容的比较。此步骤对于验证更改是否符合预期至关重要。

import os
import shutil
import zipfile
from datetime import datetime



# 定义标签目录的路径
# 用实际的标签目录路径替换
labels_directory = r"D:\Desktop\20231222\train\labels_backup_20231229152659"

# 定义备份 zip 文件路径为原目录的同级目录，带有日期
timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
backup_zip_path = os.path.join(
    os.path.dirname(labels_directory), f"labels_backup_{timestamp}.zip"
)

# 创建一个备份 zip 文件
with zipfile.ZipFile(backup_zip_path, "w") as backup_zip:
    for foldername, subfolders, filenames in os.walk(labels_directory):
        for filename in filenames:
            file_path = os.path.join(foldername, filename)
            backup_zip.write(file_path, os.path.relpath(file_path, labels_directory))

# 定义当前（旧）类别名称和新类别名称
current_names = ["blue", "green", "red", "yellow"]
new_names = ["red", "yellow", "green", "blue"]

# 创建一个从旧类别索引到新类别索引的映射，基于名称
name_to_index = {name: index for index, name in enumerate(current_names)}
index_mapping = {name_to_index[name]: new_names.index(name) for name in current_names}


# 更新标签文件中的类别索引的函数
def update_class_index(file_path, mapping):
    with open(file_path, "r") as file:
        lines = file.readlines()
    old_content = "".join(lines)
    new_lines = []
    for line in lines:
        parts = line.strip().split()
        if parts:
            class_index = int(parts[0])
            # 使用提供的映射映射类别索引
            parts[0] = str(mapping.get(class_index, class_index))
            new_lines.append(" ".join(parts))
    new_content = "\n".join(new_lines)
    with open(file_path, "w") as file:
        file.write(new_content)
    return old_content, new_content


# 记录修改的文件名称和数量
modified_file_names = []
modified_file_count = 0

# 找到最后一个修改的文件和内容
last_file_name = None
last_file_old_content = None
last_file_new_content = None

# 记录被忽略的文件后缀
ignored_file_extensions = set()

# 记录原本的文件数量
original_file_count = 0

# 将映射应用于标签目录中的所有 .txt 文件
for filename in sorted(os.listdir(labels_directory)):
    file_path = os.path.join(labels_directory, filename)

    # 过滤非文本文件
    if not filename.endswith(".txt"):
        ignored_file_extensions.add(os.path.splitext(filename)[1])
        continue

    original_file_count += 1

    old_content, new_content = update_class_index(file_path, index_mapping)
    modified_file_names.append(filename)
    modified_file_count += 1
    last_file_name = filename
    last_file_old_content = old_content
    last_file_new_content = new_content

# 输出原本的文件数量
print(f"Original number of files: {original_file_count}")

# 输出修改的文件名称和总数
print(f"Modified file names: {modified_file_names}")
print(f"Total number of files modified: {modified_file_count}")

# 打印最后一个文件的修改前后内容
if last_file_name:
    print(f"\nLast modified file: {last_file_name}")
    print("Before modification:")
    print(last_file_old_content)
    print("\nAfter modification:")
    print(last_file_new_content)

# 输出新旧类别名称和索引映射
print("\n映射结果：")
for old_index, new_name in index_mapping.items():
    old_name = current_names[old_index]
    print(
        f"Class '{old_name}' (old, index {old_index}) -> Class '{new_name}' (new, index {new_name})"
    )