当前位置：首页 > article >正文

按类别调整目标检测标注框的写入顺序以优化人工审核效率

article 2025/1/21 12:10:21

引言

在目标检测数据标注审核过程中，我们常常会遇到以下情况：某些小目标的检测框嵌套在大目标检测框内，而在模型进行预标注后，这些小目标的框可能被写入到了大目标框的下层。在人工审核阶段，标注审核人员需要手动移动或删除上层的大目标框以调整小目标框的位置，导致效率低下。

为了解决这一问题，我们可以通过自动化调整标注框的写入顺序，使小目标的检测框优先显示在上层，提升人工审核的效率。

本文将通过一个 Python 脚本展示如何实现这一功能，并对代码逻辑进行讲解和优化。

问题描述

在 XML 格式的标注文件中，标注信息是以 object 标签的形式存储的。每个 object 标签包含了目标类别、位置信息等内容。标注框的显示顺序通常与 XML 文件中标签的写入顺序一致。因此，调整标注框的写入顺序即可改变显示顺序。

为了解决小目标框被覆盖的问题，我们需要：

确定需要调整顺序的类别（如小目标类别）。
将这些类别的标注框写入到 XML 文件的后面，从而显示在上层。

代码实现

import os
import xml.etree.ElementTree as ET
from tqdm import tqdm
import xml.dom.minidom as minidom


def reorder_annotation_boxes(xml_folder: str, prioritized_categories: list):
    """
    调整目标检测标注框的顺序，使指定类别的标注框显示在上层。

    参数:
        xml_folder (str): 存储标注文件（XML格式）的文件夹路径。
        prioritized_categories (list): 需要优先显示（后写入）的目标类别列表。
    """
    # 获取文件夹中所有 XML 文件
    xml_files = [file for file in os.listdir(xml_folder) if file.endswith('.xml')]
    
    for xml_file in tqdm(xml_files, desc="Processing XML files"):
        xml_path = os.path.join(xml_folder, xml_file)
        
        # 解析 XML 文件
        tree = ET.parse(xml_path)
        root = tree.getroot()
        
        # 获取所有目标标注框
        all_objects = root.findall('object')
        
        # 按类别分类标注框
        prioritized_boxes = [obj for obj in all_objects if obj.find('name').text in prioritized_categories]
        other_boxes = [obj for obj in all_objects if obj.find('name').text not in prioritized_categories]
        
        # 清空原来的标注框
        for obj in all_objects:
            root.remove(obj)
        
        # 按顺序重新写入标注框
        for box in other_boxes:  # 非优先类别的框先写入
            root.append(box)
        for box in prioritized_boxes:  # 优先类别的框后写入
            root.append(box)
        
        # 使用 minidom 美化并保存 XML 文件
        xml_str = ET.tostring(root, encoding='utf-8')
        pretty_xml = minidom.parseString(xml_str).toprettyxml(indent="\t")
        with open(xml_path, 'w', encoding='utf-8') as output_file:
            output_file.write(pretty_xml)


if __name__ == "__main__":
    # 输入文件夹路径
    xml_folder_path = '/data/inxml'
    
    # 需要优先显示的类别列表
    prioritized_categories = [
        'cls_1', 'cls_2'
    ]
    
    # 调用函数处理 XML 文件
    reorder_annotation_boxes(xml_folder_path, prioritized_categories)