如何在TensorFlow目标检测API中为不同类别设置不同阈值可视化检测框与标签

方案1：预处理过滤检测结果（无需修改原函数，优先推荐） #

原函数的输入boxes、classes、scores是一一对应的数组，你可以先根据类别自定义阈值过滤掉不符合要求的检测结果，再把过滤后的内容传给原可视化函数即可，操作步骤如下：

• 先定义类别阈值映射字典，比如class_threshold = {1:0.9, 2:0.6, 3:0.2}
• 遍历所有检测结果，保留满足「当前类别对应阈值 ≤ 检测得分」的条目
• 把过滤后的boxes、classes、scores数组输入原visualize_boxes_and_labels_on_image_array函数，原函数的min_score_thresh直接设为0即可（已提前完成过滤，不需要二次校验）

示例代码片段：

import numpy as np

# 自定义类别阈值
class_thresh = {
    1: 0.9,
    2: 0.6,
    3: 0.2
}

# 结果过滤逻辑
keep_indices = []
for idx in range(len(scores)):
    current_class = int(classes[idx])
    current_score = scores[idx]
    # 未自定义阈值的类别默认用0.5作为阈值
    if current_score >= class_thresh.get(current_class, 0.5):
        keep_indices.append(idx)

# 提取过滤后的结果
filtered_boxes = boxes[keep_indices]
filtered_classes = classes[keep_indices]
filtered_scores = scores[keep_indices]
# 如果用到instance_masks等其他参数，也按相同索引过滤

# 调用原可视化函数
visualize_boxes_and_labels_on_image_array(
    image=image,
    boxes=filtered_boxes,
    classes=filtered_classes,
    scores=filtered_scores,
    category_index=category_index,
    min_score_thresh=0,
    # 其他参数按原有逻辑传入即可
)

方案2：修改原可视化函数内部逻辑（适合需要频繁复用分类别阈值能力的场景） #

你可以在原函数的参数列表里新增class_wise_thresh参数（默认值设为None），然后修改原函数内部的阈值判断逻辑：

• 如果class_wise_thresh为None，就沿用原来的统一min_score_thresh做判断
• 如果传入了类别阈值字典，就按每个类别的自定义阈值判断是否需要绘制当前检测框

修改后的函数关键逻辑示例：

def visualize_boxes_and_labels_on_image_array(
    # 原有参数保持不变
    min_score_thresh=.5,
    class_wise_thresh=None, # 新增自定义类别阈值参数
    # 其余原有参数保持不变
):
    # 原有函数逻辑保持不变...
    # 找到原函数中判断得分是否达标、是否要绘制框的代码位置，将原有判断逻辑替换为：
    current_thresh = min_score_thresh
    if class_wise_thresh is not None:
        current_thresh = class_wise_thresh.get(int(class_id), min_score_thresh)
    if score >= current_thresh:
        # 执行原有绘制逻辑
    # 其余函数逻辑保持不变

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