如何通过TensorFlow Object Detection API输出检测框的XML坐标文件？

嘿，我来帮你搞定这个需求！要把检测出的边界框导出成和labelImg格式一致的XML文件，其实只需要两步：获取检测结果的坐标数据，然后按照PASCAL VOC格式生成XML文件。下面是具体的实现步骤：

1. 从推理结果中提取边界框坐标 #

首先，在你的推理脚本里，TensorFlow Object Detection API的detect_fn会返回一个包含检测结果的字典detections，里面关键的字段有：

• detection_boxes：归一化的边界框坐标，格式为[y_min, x_min, y_max, x_max]（数值范围0-1）
• detection_classes：检测到的物体类别ID（注意：通常从1开始计数）
• detection_scores：每个检测结果的置信度分数

你需要把归一化坐标转换成图像的像素坐标，公式很简单：

# 假设image是你输入的图像，image_height和image_width是图像的像素尺寸
image_height, image_width = image.shape[:2]

# 遍历每个检测结果
for i in range(len(detections['detection_scores'])):
    # 先过滤低置信度的结果，比如只保留分数≥0.5的
    if detections['detection_scores'][i] < 0.5:
        continue
    # 提取归一化坐标
    y_min, x_min, y_max, x_max = detections['detection_boxes'][i]
    # 转换为像素坐标
    x_min_pix = int(x_min * image_width)
    y_min_pix = int(y_min * image_height)
    x_max_pix = int(x_max * image_width)
    y_max_pix = int(y_max * image_height)

2. 编写函数生成PASCAL VOC格式的XML #

labelImg生成的是PASCAL VOC标准的XML，我们可以用Python的xml.etree.ElementTree库来构建这个结构。下面是一个完整的生成函数：

import os
import xml.etree.ElementTree as ET

def create_voc_xml(image_path, image_width, image_height, detections, class_names, output_path):
    # 创建XML根节点
    root = ET.Element("annotation")
    
    # 添加folder节点
    folder = ET.SubElement(root, "folder")
    folder.text = os.path.basename(os.path.dirname(image_path))
    
    # 添加filename节点
    filename = ET.SubElement(root, "filename")
    filename.text = os.path.basename(image_path)
    
    # 添加path节点
    path = ET.SubElement(root, "path")
    path.text = image_path
    
    # 添加source节点
    source = ET.SubElement(root, "source")
    database = ET.SubElement(source, "database")
    database.text = "Unknown"
    
    # 添加size节点（图像尺寸）
    size = ET.SubElement(root, "size")
    width_elem = ET.SubElement(size, "width")
    width_elem.text = str(image_width)
    height_elem = ET.SubElement(size, "height")
    height_elem.text = str(image_height)
    depth_elem = ET.SubElement(size, "depth")
    depth_elem.text = "3"  # 彩色图像是3，灰度图改成1
    
    # 添加segmented节点（固定为0即可）
    segmented = ET.SubElement(root, "segmented")
    segmented.text = "0"
    
    # 遍历所有符合条件的检测结果，添加object节点
    for i in range(len(detections['detection_scores'])):
        score = detections['detection_scores'][i]
        # 过滤低置信度结果，可根据需求调整阈值
        if score < 0.5:
            continue
        
        # 获取类别名称（注意class_id从1开始，所以要减1对应列表索引）
        class_id = int(detections['detection_classes'][i])
        class_name = class_names[class_id - 1]
        
        # 转换坐标为像素值
        y_min, x_min, y_max, x_max = detections['detection_boxes'][i]
        x_min_pix = int(x_min * image_width)
        y_min_pix = int(y_min * image_height)
        x_max_pix = int(x_max * image_width)
        y_max_pix = int(y_max * image_height)
        
        # 创建object节点
        obj = ET.SubElement(root, "object")
        name_elem = ET.SubElement(obj, "name")
        name_elem.text = class_name
        pose_elem = ET.SubElement(obj, "pose")
        pose_elem.text = "Unspecified"
        truncated_elem = ET.SubElement(obj, "truncated")
        truncated_elem.text = "0"
        difficult_elem = ET.SubElement(obj, "difficult")
        difficult_elem.text = "0"
        
        # 添加边界框节点
        bndbox = ET.SubElement(obj, "bndbox")
        xmin_elem = ET.SubElement(bndbox, "xmin")
        xmin_elem.text = str(x_min_pix)
        ymin_elem = ET.SubElement(bndbox, "ymin")
        ymin_elem.text = str(y_min_pix)
        xmax_elem = ET.SubElement(bndbox, "xmax")
        xmax_elem.text = str(x_max_pix)
        ymax_elem = ET.SubElement(bndbox, "ymax")
        ymax_elem.text = str(y_max_pix)
    
    # 生成XML文件并保存
    tree = ET.ElementTree(root)
    # 写入时指定编码和XML声明，和labelImg的输出保持一致
    tree.write(output_path, encoding='utf-8', xml_declaration=True)

3. 在推理代码中调用函数 #

把上面的函数整合到你的推理流程里，示例代码如下：

# 假设你已经加载了检测模型，定义了detect_fn
import cv2

# 输入图像路径
image_path = "test_image.jpg"
# 读取图像
image = cv2.imread(image_path)
image_height, image_width = image.shape[:2]
# 执行推理
detections = detect_fn(image)
# 你的类别名称列表，要和训练时的label map完全对应
class_names = ["cat", "dog", "bird"]  # 示例类别，替换成你的实际类别
# 输出XML文件的路径
output_xml_path = "test_image_output.xml"

# 调用函数生成XML
create_voc_xml(image_path, image_width, image_height, detections, class_names, output_xml_path)

这样生成的XML文件就和你用labelImg标注的格式完全一致了！如果需要调整置信度阈值、修改XML里的字段（比如database内容），直接修改函数里的对应部分就行。

内容的提问来源于stack exchange，提问作者Ajinkya