如何从数量未知的边界框创建TensorFlow图像掩码张量？

嗨，这个问题我之前也碰到过！当边界框数量不确定时，确实没法用tf.unstack逐个处理，但我们可以用TensorFlow的向量化操作高效实现，完全不需要循环或者提前框的数量。下面是具体的思路和代码实现：

核心思路 #

我们可以生成整个图像的像素坐标网格，然后一次性判断每个像素是否落在任意一个边界框内，最后生成对应的掩码。这种方法依赖TensorFlow的广播机制，效率很高，还能适配任意数量的边界框。

完整代码实现 #

假设你的边界框格式是[N, 4]，每个框的格式为(y_min, x_min, y_max, x_max)（绝对像素坐标，左上角为原点），图像形状为(height, width, channels)：

import tensorflow as tf

def create_mask_from_bounding_boxes(image, boxes):
    # 获取图像的高度、宽度
    height = tf.shape(image)[0]
    width = tf.shape(image)[1]
    
    # 生成像素坐标网格：(height, width, 2)，每个元素是(y, x)
    y_coords, x_coords = tf.meshgrid(tf.range(height), tf.range(width), indexing='ij')
    coords = tf.stack([y_coords, x_coords], axis=-1)
    coords = tf.expand_dims(coords, axis=2)  # 形状变为(height, width, 1, 2)，方便和boxes广播
    
    # 处理边界框：将boxes调整为(1, 1, N, 4)，拆分出四个坐标值
    boxes = tf.expand_dims(tf.expand_dims(boxes, axis=0), axis=0)
    y_min, x_min, y_max, x_max = tf.split(boxes, 4, axis=-1)
    
    # 判断每个像素是否在任意一个边界框内
    in_box = tf.logical_and(
        tf.logical_and(coords[..., 0] >= y_min, coords[..., 0] < y_max),
        tf.logical_and(coords[..., 1] >= x_min, coords[..., 1] < x_max)
    )
    # 对所有边界框取逻辑或，得到每个像素是否在任意框内的掩码：(height, width)
    mask = tf.reduce_any(in_box, axis=-1)
    
    # 将掩码扩展到通道维度，方便和图像相乘：(height, width, 1)
    mask = tf.expand_dims(mask, axis=-1)
    
    # 生成填充零后的图像：框内的区域变为0，其余保留原图像
    masked_image = image * tf.cast(1 - mask, dtype=image.dtype)
    
    return masked_image, mask

关键步骤解释 #

1. 坐标网格生成：用tf.meshgrid生成每个像素的(y, x)坐标，调整形状后能和任意数量的边界框进行广播比较。
2. 广播比较：把边界框的形状扩展为(1,1,N,4)，坐标网格扩展为(height,width,1,2)，TensorFlow会自动广播，一次性计算所有像素和所有边界框的位置关系。
3. 逻辑或聚合：用tf.reduce_any对所有边界框的判断结果取逻辑或，得到每个像素是否在任意框内的掩码。
4. 掩码应用：将掩码转换为和图像同类型的张量，乘以1 - mask后和原图像相乘，就能把框内的区域置零。

适配归一化边界框 #

如果你的边界框是归一化坐标（值在0-1之间），只需要在处理前把框转换为绝对坐标即可：

# 转换归一化框为绝对坐标
boxes = tf.stack([
    boxes[..., 0] * tf.cast(height, tf.float32),
    boxes[..., 1] * tf.cast(width, tf.float32),
    boxes[..., 2] * tf.cast(height, tf.float32),
    boxes[..., 3] * tf.cast(width, tf.float32)
], axis=-1)
boxes = tf.cast(boxes, tf.int32)  # 如果需要整数坐标

这种方法完全不需要提前知道边界框的数量，而且所有操作都是TensorFlow的图操作，支持自动微分和GPU部署，效率拉满！

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