Tensorflow keras中实现语义分割多分类指标：IOU、MIOU

在TF1.x版本中 miou指标可以使用tf.metrics.mean_iou 进行计算：

tf.metrics.mean_iou(labels, predictions, num_classes) 

但是该方法有如下几点限制：

1. 无法在动态图中使用，例如Tensorflow2.x版本中（注：TF2.x中api移动到了tf.compat.v1.metrics.mean_iou中），由于TF2.x默认是开启动态图，因此会报错（见mean_iou方法的源码）

  if context.executing_eagerly():
    raise RuntimeError('tf.metrics.mean_iou is not supported when '
                       'eager execution is enabled.')

2. 使用必须先 sess.run(tf.local_variables_initializer()) 然后 sess.run(update_op)，最后sess.run(mean_iou_v)，注意次序不能颠倒，不太方便和tf.keras相关训练代码结合使用

mean_iou_v, update_op = tf.metrics.mean_iou(y_true, y_pred, num_classes=4)
sess = tf.Session()
sess.run(tf.local_variables_initializer())
print(sess.run(update_op))
print(sess.run(mean_iou_v))

3. 只能直接输出所有类别的平均IOU即mean_iou, 而不能输出各个类别对应的 iou

针对上述三个问题，我发现有如下两种解决方案：

目录

方案1：自己实现相关计算代码

方案2：继承调用tf.keras.metrics.MeanIoU类

---

方案1：自己实现相关计算代码

def cal_mean_iou(num_classes, ignore_labels=None):
    """
    num_classes: int, 表示类别总数
    ignore_labels: list[int]，注意这里ignore_labels必须为列表或None，
    若为列表则存放int类型数字，表示需要忽略（不需要计算miou）的类别，
    例如：num_classes=12 ignore_labels=[11] 表示总类别数为12，忽略第11个类别
    """
    
    def MIOU(y_true, y_pred):
        """
        y_true: Tensor，真实标签（one-hot类型），
        y_pred: Tensor，模型输出结果（one-hot类型），二者shape都为[N,H,W,C]或[N,H*W,C],C为总类别数,
        """
        y_true = tf.reshape(tf.argmax(y_true, axis=-1), [-1])  # 求argmax后，展平为一维
        y_pred = tf.reshape(tf.argmax(y_pred, axis=-1), [-1])
        num_need_labels = num_classes #不忽略的类别总数
        if ignore_labels is not None:
            num_need_labels -= len(ignore_labels)
            for ig in ignore_labels:
                mask = tf.not_equal(y_true, ignore_labels)  # 获取需要忽略的标签的位置
                y_true = tf.boolean_mask(y_true, mask)  # 剔除y_true中需要忽略的标签
                y_pred = tf.boolean_mask(y_pred, mask)  # 剔除y_pred中需要忽略的标签
        confusion_matrix = tf.confusion_matrix(y_true, y_pred, num_classes)  # 计算混淆矩阵
        intersect = tf.diag_part(confusion_matrix)  # 获取对角线上的矩阵，形成一维向量
        union = tf.reduce_sum(confusion_matrix, axis=0) + tf.reduce_sum(confusion_matrix, axis=1) - intersect
        iou = tf.div_no_nan(tf.cast(intersect, tf.float32), tf.cast(union, tf.float32))
        num_valid_entries = tf.reduce_sum(tf.cast(tf.not_equal(union, 0), dtype=tf.float32)) #统计union中不为0的总数
        num = tf.minimum(num_valid_entries, num_need_labels)
        mean_iou = tf.div_no_nan(tf.reduce_sum(iou), num)  # mean_iou只需要计算union中不为0且不忽略label的
        return mean_iou

    return MIOU

上述代码是自己实现的各类别IOU以及平均IOU的计算方法，

如果只是想直接显示平均IOU，那么直接这样使用即可:

model.compile(optimizer=optimizer, 
              loss=loss,
              metrics=[cal_mean_iou(num_classes, ignore_label)])

如果想要在tf.keras训练过程中显示各个类别的IOU，一般是继承tf.keras.callbacks.Callback类，然后重写相关的方法，方法可参考：相关参考博客3

方案2：继承调用tf.keras.metrics.MeanIoU类

方案1中的计算方式和tf.keras.metrics.MeanIoU(num_classes)计算方式类似，需要注意tf.keras.metrics.MeanIoU类中update_state(self, y_true, y_pred, sample_weight=None)方法接受的y_true和y_pred一般是非one-hot编码形式的，即如果网络的输入shape为[N,H,W,C]或[N,H*W,C]形式，需要将y_true和y_pred先求argmax，然后调用该方法

因此遇到上述情况，可以先继承tf.keras.metrics.MeanIoU类，然后重写update_state方法，示例代码如下：

class MeanIoU(tf.keras.metrics.MeanIoU):
    """
    y_true: Tensor，真实标签（one-hot类型），
    y_pred: Tensor，模型输出结果（one-hot类型），二者shape都为[N,H,W,C],C为总类别数,
    """
    def update_state(self, y_true, y_pred, sample_weight=None):
        y_true = tf.argmax(y_true, axis=-1)
        y_pred = tf.argmax(y_pred, axis=-1)
        super().update_state(y_true, y_pred, sample_weight=sample_weight)
        return self.result()

model.compile(optimizer=optimizer, 
              loss=loss,
              metrics=[MeanIoU(num_classes)])

相关参考博客：

1.图像分割常用指标及MIoU计算 - 简书

2.MIoU 源码解析 - Wenshan's Blog

3.Keras上实现recall和precision，f1-score（多分类问题）_热爱学习的Valeria的博客-CSDN博客_keras precision

4.Tensorflow中tf.keras.metrics.MeanIoU在shape不一致错误_Bluish White的博客-CSDN博客