如何在Tensorflow中正确使用tf.metrics.mean_iou来显示Tensorboard上的混淆矩阵？

import tensorflow as tf
from tensorflow.python.ops import control_flow_ops

def compute_confusion_matrix(labels, predictions, num_classes):
    # 将预测结果和真实标签转换为一维向量
    predictions = tf.reshape(predictions, [-1])
    labels = tf.reshape(labels, [-1])

    # 创建混淆矩阵
    confusion_matrix = tf.confusion_matrix(labels, predictions, num_classes=num_classes)

    return confusion_matrix

# 定义标签和预测结果
labels = tf.placeholder(tf.int32, shape=[None, height, width, num_classes], name='labels')
predictions = tf.placeholder(tf.int32, shape=[None, height, width, num_classes], name='predictions')

# 计算IoU指标
mean_iou, update_op = tf.metrics.mean_iou(labels, predictions, num_classes=num_classes)

# 获取混淆矩阵
confusion_matrix = compute_confusion_matrix(labels, predictions, num_classes=num_classes)

# 添加混淆矩阵到Tensorboard
tf.summary.image('Confusion Matrix', tf.expand_dims(tf.cast(confusion_matrix, tf.float32), axis=0))

# 合并所有的summary
summary_op = tf.summary.merge_all()

with tf.Session() as sess:
    # 初始化变量
    sess.run(tf.global_variables_initializer())
    sess.run(tf.local_variables_initializer())

    # 创建summary writer
    writer = tf.summary.FileWriter(logdir)

    # 训练过程中更新summary和混淆矩阵
    for step in range(num_steps):
        # 执行训练操作
        sess.run(train_op)

        # 更新混淆矩阵和IoU指标
        sess.run(update_op, feed_dict={labels: labels_batch, predictions: predictions_batch})

        # 每隔一定步数写入summary
        if step % summary_interval == 0:
            summary = sess.run(summary_op, feed_dict={labels: labels_batch, predictions: predictions_batch})
            writer.add_summary(summary, global_step=step)

    writer.close()