基于标准TensorFlow算子从掩码图像生成边界框

核心方案：基于tf.image.find_contours的多边界框生成 #

使用TensorFlow核心算子实现，完全支持TFLite转换，无需依赖第三方库或自定义算子，步骤如下：

代码实现 #

import tensorflow as tf
import cv2
import os

# 加载并生成布尔掩码
mask = cv2.imread(os.path.expanduser('~/Downloads/example_mask3.png')).mean(axis=2) > 50
# 转换为find_contours要求的uint8格式
mask_uint8 = tf.cast(mask, tf.uint8) * 255

# 提取掩码中所有真值区域的轮廓
contours = tf.image.find_contours(mask_uint8, threshold=127)

# 为每个轮廓计算最小包围边界框
bounding_boxes = []
for contour in contours:
    x_coords = contour[:, 1]
    y_coords = contour[:, 0]
    xmin = tf.reduce_min(x_coords)
    ymin = tf.reduce_min(y_coords)
    xmax = tf.reduce_max(x_coords)
    ymax = tf.reduce_max(y_coords)
    bounding_boxes.append(tf.stack([xmin, ymin, xmax, ymax]))

# 转换为张量格式（可选，适配模型处理流程）
bounding_boxes_tensor = tf.stack(bounding_boxes)

方案优势 #

• 所有操作均使用TensorFlow核心算子，无tfa.image.connected_components这类非TFLite兼容算子；
• tf.image.find_contours直接提取每个连通真值区域的轮廓，生成的边界框精准对应目标区域；
• 可直接嵌入TensorFlow模型中，转换为TFLite时无需额外编译或自定义算子支持。

可选优化与替代方案 #

1. 噪点预处理：如果掩码存在小噪点，可先用tf.image.erode_dilate做形态学滤波，去除孤立噪点后再提取轮廓：

   # 定义3x3的结构元素
   kernel = tf.ones((3,3), dtype=tf.uint8)
   # 先腐蚀再膨胀去噪
   mask_cleaned = tf.image.erode_dilate(mask_uint8, kernel, iterations=1)
   

2. 非TensorFlow环境替代：如果仅在预处理阶段使用，可采用OpenCV的连通组件分析：

   import numpy as np
   _, labels = cv2.connectedComponents(mask_uint8.numpy())
   bounding_boxes = []
   for label in range(1, labels.max()+1):
       y, x = np.where(labels == label)
       bounding_boxes.append([x.min(), y.min(), x.max(), y.max()])
   

   注：此方案无法嵌入TensorFlow模型，仅适用于离线预处理。

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