TensorFlow目标检测API：如何获取分类标签与预测概率

嘿，我来帮你理清楚怎么从TensorFlow目标检测的结果里提取分类标签和预测概率～结合你提供的代码片段和官方教程的逻辑，我给你拆解成具体实操步骤：

1. 先明确检测结果的核心结构 #

在TensorFlow Object Detection API中，模型推理后会返回几个关键张量：

• detection_classes：每个检测框对应的分类ID（返回值是浮点型，需要转成整数才能用）
• detection_scores：每个检测框对应的预测概率（范围0-1，数值越高置信度越强）
• detection_boxes：你已经提到的检测框坐标信息

2. 加载标签映射文件 #

官方教程里一定会用到label_map.pbtxt文件，它是分类ID和可读标签文本的对应表。先把它加载成可直接查询的字典：

from object_detection.utils import label_map_util

# 替换成你自己的label_map文件路径
label_map_path = 'path/to/your/label_map.pbtxt'
label_map = label_map_util.load_labelmap(label_map_path)
categories = label_map_util.convert_label_map_to_categories(
    label_map, max_num_classes=90, use_display_name=True)
# 生成ID到标签的映射字典，方便后续查询
category_index = label_map_util.create_category_index(categories)

3. 嵌入你的图片循环提取结果 #

把这段逻辑加到你现有的代码里，就能直接拿到每张图片的标签和概率了：

for image_path in TEST_IMAGE_PATHS:
    image = Image.open(image_path)
    image_np = np.array(image)
    # 扩展维度，匹配模型的输入格式要求
    image_np_expanded = np.expand_dims(image_np, axis=0)
    
    # 执行模型推理，获取所有检测结果
    (boxes, scores, classes, num_detections) = sess.run(
        [detection_boxes, detection_scores, detection_classes, num_detections],
        feed_dict={image_tensor: image_np_expanded})
    
    # 遍历有效检测结果（建议过滤低置信度结果，比如只保留>0.5的）
    for i in range(int(num_detections[0])):
        # 把浮点型的分类ID转成整数，才能匹配映射字典的键
        class_id = int(classes[0][i])
        # 获取当前检测框的预测概率
        confidence_score = scores[0][i]
        # 通过ID映射到对应的可读标签文本
        class_label = category_index[class_id]['name']
        
        # 这里可以根据需求处理结果，比如打印或用于可视化
        print(f"检测到目标: {class_label}, 置信度: {confidence_score:.2f}")

几个关键提醒 #

• 必须把classes里的浮点型转成整数，不然无法匹配category_index的键值对
• 记得设置置信度阈值（比如confidence_score > 0.5），过滤掉低概率的误检结果
• 如果是自定义训练的模型，要保证label_map.pbtxt和训练时使用的完全一致，否则标签会对应错误

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