如何将TensorFlow Lite的C++输出转换为图像分割掩码？

方案一：在C++中对float输出做后处理 #

TF-Lite输出的float张量和Python中model.predict的输出完全一致，只需复刻Python里的二值化逻辑即可：

1. 获取TF-Lite输出张量 #

假设你已经完成模型加载和推理，先取出输出张量的float数据：

// 已初始化并运行的TF-Lite解释器
TfLiteInterpreter* interpreter = ...;
// 获取第一个输出张量（对应模型的掩码输出）
TfLiteTensor* output_tensor = interpreter->output_tensor(0);
// 输出张量形状为[1, 128, 128, 1]，直接读取float数据
float* output_data = output_tensor->data.f;

2. 二值化转换 #

遍历每个像素值，复刻Python中>=0.5则设为255，否则0的逻辑：

// 存储最终掩码的数组，大小为128*128
uint8_t mask_pixels[128 * 128];
int pixel_idx = 0;

// 按行遍历图像（对应Python中的y轴）
for (int y = 0; y < 128; ++y) {
    for (int x = 0; x < 128; ++x) {
        // 从输出张量中取出对应像素的float值
        float pixel_val = output_data[y * 128 + x];
        // 二值化转换
        mask_pixels[pixel_idx++] = (pixel_val >= 0.5f) ? 255 : 0;
    }
}

3. 生成图像 #

如果需要保存为图像文件，可使用OpenCV等库实现：

// 用OpenCV创建单通道灰度图
cv::Mat mask_img(128, 128, CV_8UC1, mask_pixels);
// 保存为PNG格式
cv::imwrite("binary_mask.png", mask_img);

方案二：转换TF-Lite时嵌入二值化层（推荐） #

可以在Python转换模型阶段，把二值化逻辑直接加入模型，让TF-Lite输出直接是uint8类型的0/255掩码，省去C++后处理步骤：

1. 修改原模型添加二值化层 #

import tensorflow as tf

# 假设原图像分割模型为model
# 定义二值化操作：复刻Python中的后处理逻辑
def binarize_output(x):
    # 先判断是否>=0.5，转换为uint8类型，再乘以255得到0/255的掩码
    x = tf.cast(tf.greater_equal(x, 0.5), tf.uint8)
    x = x * 255
    return x

# 创建包含二值化层的新模型
new_model = tf.keras.Model(inputs=model.input, outputs=binarize_output(model.output))

2. 转换为优化后的TF-Lite模型 #

# 初始化转换器
converter = tf.lite.TFLiteConverter.from_keras_model(new_model)
# 开启默认优化
converter.optimizations = [tf.lite.Optimize.DEFAULT]
# 指定输出类型为uint8
converter.target_spec.supported_types = [tf.uint8]

# 转换并保存模型
tflite_model = converter.convert()
with open("segmentation_model_binarized.tflite", "wb") as f:
    f.write(tflite_model)

3. C++中直接使用输出 #

加载这个模型后，输出张量的类型为uint8，直接读取数据即可作为图像使用：

TfLiteTensor* output_tensor = interpreter->output_tensor(0);
uint8_t* mask_data = output_tensor->data.uint8;
// 直接用mask_data创建图像，无需额外转换
cv::Mat mask_img(128, 128, CV_8UC1, mask_data);

注意事项 #

• 确保C++中模型输入的预处理逻辑和Python完全一致（比如灰度图的归一化、输入张量的形状顺序等），否则输出会出现偏差。
• 若输出张量形状与预期不符，可通过output_tensor->dims查看维度信息，调整遍历逻辑。

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