如何从Keras向TensorBoard写入图像？分割任务预测可视化问询

我来帮你搞定这个问题！要在Keras的分割任务里把预测结果可视化到TensorBoard，核心思路是自定义一个回调函数——毕竟默认的TensorBoard回调只处理标量指标，不支持图像类的可视化。下面给你一套实用的方案：

1. 导入所需依赖 #

首先要导入TensorFlow和Keras的相关模块：

import tensorflow as tf
from tensorflow.keras.callbacks import Callback
import numpy as np

2. 自定义图像可视化回调类 #

这个回调会在每个epoch结束时，从验证集中抽取样本，运行模型预测，然后把输入图像、真实标签、预测结果拼接起来写入TensorBoard：

class SegmentationVisCallback(Callback):
    def __init__(self, val_data, log_dir='./segmentation_logs', num_samples=5):
        super().__init__()
        self.val_data = val_data  # 你的验证数据集
        self.num_samples = num_samples  # 每次要可视化的样本数量
        self.writer = tf.summary.create_file_writer(log_dir)  # 创建Summary写入器

    def on_epoch_end(self, epoch, logs=None):
        # 遍历要可视化的样本
        for idx in range(self.num_samples):
            # 获取单个样本（区分两种验证集格式）
            if isinstance(self.val_data, tf.data.Dataset):
                # 如果是tf.data.Dataset格式
                x_batch, y_true_batch = next(iter(self.val_data))
                x = x_batch[idx:idx+1]
                y_true = y_true_batch[idx:idx+1]
            else:
                # 如果是numpy数组格式（(x_val, y_val)）
                x = self.val_data[0][idx:idx+1]
                y_true = self.val_data[1][idx:idx+1]
            
            # 运行模型预测
            y_pred = self.model.predict(x, verbose=0)
            
            # 处理图像格式，转成可显示的0-255灰度/彩色图
            # 输入图像（假设已归一化到0-1，转成0-255）
            x_img = tf.cast(tf.squeeze(x) * 255, tf.uint8)
            # 真实标签（假设标签是0-1的概率值，转成0-255）
            y_true_img = tf.cast(tf.squeeze(y_true) * 255, tf.uint8)
            # 预测结果（同样处理）
            y_pred_img = tf.cast(tf.squeeze(y_pred) * 255, tf.uint8)
            
            # 把三张图横向拼接，方便对比
            combined_img = tf.concat([x_img, y_true_img, y_pred_img], axis=1)
            # 增加通道维度（TensorBoard要求图像格式是[batch, height, width, channels]）
            combined_img = tf.expand_dims(combined_img, axis=-1)
            
            # 写入TensorBoard
            with self.writer.as_default():
                tf.summary.image(f"Segmentation_Sample_{idx}", combined_img, step=epoch)
        
        # 刷新写入器，确保数据保存
        self.writer.flush()

3. 在训练中使用回调 #

把自定义回调和默认的TensorBoard回调一起加入训练流程即可：

from tensorflow.keras.callbacks import TensorBoard

# 假设你已经准备好训练集train_dataset、验证集val_dataset和模型model
log_dir = './segmentation_logs'

# 初始化自定义回调和TensorBoard回调
seg_callback = SegmentationVisCallback(val_data=val_dataset, log_dir=log_dir, num_samples=5)
tb_callback = TensorBoard(log_dir=log_dir)

# 开始训练
model.fit(
    train_dataset,
    validation_data=val_dataset,
    epochs=30,
    callbacks=[tb_callback, seg_callback]
)

4. 关键注意事项 #

• 多类分割适配：如果是多类别分割，需要把预测的概率图转成类别标签，比如用tf.argmax(y_pred, axis=-1)获取类别索引，再按类别数缩放灰度值：

  num_classes = 3  # 假设你的任务是3类分割
  y_pred_img = tf.cast(tf.argmax(tf.squeeze(y_pred), axis=-1) * (255/(num_classes-1)), tf.uint8)
  

• 图像格式调整：如果你的输入是RGB彩色图，不需要额外增加通道维度，直接拼接即可；如果是其他预处理方式（比如归一化到-1到1），要调整转成0-255的逻辑。
• 性能优化：如果验证集很大，不要每次取太多样本，避免拖慢训练速度；也可以固定抽取特定样本，方便对比不同epoch的变化。

启动训练后，打开TensorBoard（运行tensorboard --logdir=./segmentation_logs），在"Images"标签下就能看到每个epoch的输入、真实标签和预测结果对比图了！

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