数据集,批量更新分类数值OR批量删除分类行数据

数据集批量更新分类OR删除分类行数据

python 复制代码
import os


def remove_class_from_file(file_path, class_to_remove):
    """
    从YOLO格式的标注文件中删除指定类别的行记录,并去除空行。

    :param file_path: YOLO标注文件路径
    :param class_to_remove: 需要删除的类别ID
    :return: 删除的行数
    """
    deleted_lines = 0

    # 读取原始数据
    with open(file_path, 'r') as file:
        lines = file.readlines()

    # 过滤掉指定类别的行
    filtered_lines = [line for line in lines if not line.startswith(f"{class_to_remove} ")]

    # 去除空行
    cleaned_lines = [line for line in filtered_lines if line.strip()]

    # 计算删除的行数
    deleted_lines = len(lines) - len(cleaned_lines)

    # 将清理后的数据写回文件
    with open(file_path, 'w') as file:
        file.writelines(cleaned_lines)

    return deleted_lines


def update_class_in_file(file_path, old_class, new_class):
    """
    将YOLO格式的标注文件中指定类别的数字更新为新的类别ID。

    :param file_path: YOLO标注文件路径
    :param old_class: 需要更新的旧类别ID
    :param new_class: 更新后的新类别ID
    :return: 更新的行数
    """
    updated_lines = 0

    # 读取原始数据
    with open(file_path, 'r') as file:
        lines = file.readlines()

    # 更新类别ID
    updated_lines = 0
    new_lines = []
    for line in lines:
        if line.startswith(f"{old_class} "):
            new_lines.append(line.replace(f"{old_class} ", f"{new_class} ", 1))
            updated_lines += 1
        else:
            new_lines.append(line)

    # 将更新后的数据写回文件
    with open(file_path, 'w') as file:
        file.writelines(new_lines)

    return updated_lines


def process_directory(directory_path, class_to_remove=None, old_class=None, new_class=None):
    """
    处理指定目录下的所有YOLO标注文件,可以选择删除指定类别的记录或更新分类数字。

    :param directory_path: 目录路径
    :param class_to_remove: 需要删除的类别ID(如果需要删除指定类别的记录)
    :param old_class: 需要更新的旧类别ID(如果需要更新分类数字)
    :param new_class: 更新后的新类别ID(如果需要更新分类数字)
    """
    total_deleted_lines = 0
    total_updated_lines = 0

    # 遍历目录中的文件
    for filename in os.listdir(directory_path):
        if filename.endswith('.txt'):  # 假设标注文件都是以 .txt 结尾的
            file_path = os.path.join(directory_path, filename)

            if class_to_remove is not None:
                deleted_lines = remove_class_from_file(file_path, class_to_remove)
                print(f"已从文件 {filename} 中删除类别 {class_to_remove} 的记录,共删除 {deleted_lines} 行。")
                total_deleted_lines += deleted_lines

            if old_class is not None and new_class is not None:
                updated_lines = update_class_in_file(file_path, old_class, new_class)
                print(f"已从文件 {filename} 中将类别 {old_class} 更新为 {new_class},共更新 {updated_lines} 行。")
                total_updated_lines += updated_lines

    if class_to_remove is not None:
        print(f"总共删除了 {total_deleted_lines} 行记录。")

    if old_class is not None and new_class is not None:
        print(f"总共更新了 {total_updated_lines} 行记录。")


# 使用示例
directory_path = 'C:\\Users\\合格的牛马\\Desktop\\datasets\\test-delete\\labels'  # 替换为你的目录路径

# 删除指定类别的记录
class_to_remove = 5  # 替换为你要删除的类别ID
process_directory(directory_path, class_to_remove=class_to_remove)

# 更新指定类别的分类数字
old_class = 1  # 替换为旧类别ID
new_class = 100  # 替换为新类别ID
process_directory(directory_path, old_class=old_class, new_class=new_class)
相关推荐
春哥的研究所10 分钟前
AI人工智能名片小程序源码系统,名片小程序+分销商城+AI客服,包含完整搭建教程
人工智能·微信小程序·小程序
ahead~14 分钟前
【大模型入门】访问GPT_API实战案例
人工智能·python·gpt·大语言模型llm
喜欢吃豆14 分钟前
深入企业内部的MCP知识(三):FastMCP工具转换(Tool Transformation)全解析:从适配到增强的工具进化指南
java·前端·人工智能·大模型·github·mcp
pany21 分钟前
写代码的节奏,正在被 AI 改写
前端·人工智能·aigc
我爱一条柴ya1 小时前
【AI大模型】神经网络反向传播:核心原理与完整实现
人工智能·深度学习·神经网络·ai·ai编程
万米商云1 小时前
企业物资集采平台解决方案:跨地域、多仓库、百部门——大型企业如何用一套系统管好百万级物资?
大数据·运维·人工智能
新加坡内哥谈技术1 小时前
Google AI 刚刚开源 MCP 数据库工具箱,让 AI 代理安全高效地查询数据库
人工智能
慕婉03071 小时前
深度学习概述
人工智能·深度学习
大模型真好玩1 小时前
准确率飙升!GraphRAG如何利用知识图谱提升RAG答案质量(额外篇)——大规模文本数据下GraphRAG实战
人工智能·python·mcp
19891 小时前
【零基础学AI】第30讲:生成对抗网络(GAN)实战 - 手写数字生成
人工智能·python·深度学习·神经网络·机器学习·生成对抗网络·近邻算法