为TensorFlow/Models目标检测Demo添加检测目标计数功能

我太懂处理大规模图像时逐张核对结果的痛苦了，给脚本加个计数功能绝对能帮你节省大量时间。下面结合你提供的run_inference_for_single_image函数，一步步实现目标总数统计：

1. 修改单张图像推理函数，返回有效目标数 #

首先，我们要让推理函数不仅输出检测结果，还能返回当前图像中置信度达标的目标数量（过滤低置信度结果，避免统计误检）。修改后的函数如下：

def run_inference_for_single_image(image, graph, score_threshold=0.5):
    with graph.as_default():
        with tf.Session() as sess:
            # 保留原教程中的张量获取逻辑
            ops = tf.get_default_graph().get_operations()
            all_tensor_names = {output.name for op in ops for output in op.outputs}
            tensor_dict = {}
            for key in [
                'num_detections', 'detection_boxes', 'detection_scores',
                'detection_classes', 'detection_masks'
            ]:
                tensor_name = key + ':0'
                if tensor_name in all_tensor_names:
                    tensor_dict[key] = tf.get_default_graph().get_tensor_by_name(tensor_name)
            
            # 保留原图像输入处理逻辑
            if 'detection_masks' in tensor_dict:
                detection_boxes = tf.squeeze(tensor_dict['detection_boxes'], [0])
                detection_masks = tf.squeeze(tensor_dict['detection_masks'], [0])
                real_num_detection = tf.cast(tensor_dict['num_detections'][0], tf.int32)
                detection_boxes = tf.slice(detection_boxes, [0, 0], [real_num_detection, -1])
                detection_masks = tf.slice(detection_masks, [0, 0, 0], [real_num_detection, -1, -1])
                detection_masks_reframed = utils_ops.reframe_box_masks_to_image_masks(
                    detection_masks, detection_boxes, image.shape[0], image.shape[1])
                detection_masks_reframed = tf.cast(
                    tf.greater(detection_masks_reframed, 0.5), tf.uint8)
                tensor_dict['detection_masks'] = tf.expand_dims(
                    detection_masks_reframed, 0)
            
            # 保留原推理执行逻辑
            image_tensor = tf.get_default_graph().get_tensor_by_name('image_tensor:0')
            output_dict = sess.run(tensor_dict,
                                feed_dict={image_tensor: np.expand_dims(image, 0)})
            
            # 保留原数据类型转换逻辑
            output_dict['num_detections'] = int(output_dict['num_detections'][0])
            output_dict['detection_classes'] = output_dict[
                'detection_classes'][0].astype(np.int64)
            output_dict['detection_boxes'] = output_dict['detection_boxes'][0]
            output_dict['detection_scores'] = output_dict['detection_scores'][0]
            if 'detection_masks' in output_dict:
                output_dict['detection_masks'] = output_dict['detection_masks'][0]
            
            # 新增：统计当前图像的有效目标数
            valid_detections = sum(score >= score_threshold for score in output_dict['detection_scores'])
            
            # 返回检测结果 + 有效目标数
            return output_dict, valid_detections

关键改动说明：

• 新增score_threshold参数，可自定义过滤置信度（默认0.5，可根据模型精度调整）
• 计算valid_detections，统计分数达标的目标数量
• 函数返回值新增计数结果，不影响原有的检测输出

2. 遍历数据集时累加全局计数 #

在遍历图像数据集的代码块里，初始化全局计数器，每次处理完一张图像就累加计数：

# 初始化全局目标总数
total_objects = 0

# 假设你的图像路径存储在image_paths列表中
for idx, image_path in enumerate(image_paths):
    # 保留原图像加载逻辑
    image = Image.open(image_path)
    image_np = load_image_into_numpy_array(image)
    
    # 调用修改后的推理函数
    output_dict, current_count = run_inference_for_single_image(image_np, detection_graph)
    
    # 累加计数
    total_objects += current_count
    
    # 可选：打印处理进度，比如每100张输出一次
    if (idx + 1) % 100 == 0:
        print(f"已处理{idx+1}/{len(image_paths)}张图像，累计检测到{total_objects}个目标")

# 全部处理完成后输出最终总数
print(f"所有图像处理完成！累计检测到目标总数：{total_objects}")

额外优化建议 #

• 如果只想统计特定类别目标（比如COCO数据集里的行人），可以修改计数逻辑：

  # 示例：仅统计类别ID为1的目标
  valid_detections = sum(
      score >= score_threshold and cls == 1 
      for score, cls in zip(output_dict['detection_scores'], output_dict['detection_classes'])
  )
  

• 若数据集极大，可考虑加入多线程/多进程加速，但需注意TensorFlow会话的线程安全问题。

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