tensorflow语义分割计算mIoU时忽略某一类别

在做语义分割的时候，有不少数据集中的部分类别是我们不想让其参与mIoU计算的。tensorflow忽略某一类别计算mIoU的方法网上比较少。经过阅读tensorflow文档，总结出一下求解方法：

数据准备：

1、先理解一维张量：

其实数组和张量相通，tensorflow绝大多数情况下能自动转化

假设有张量a，b

a = np.array([0, 1, 2, 3])    # 真实
b = np.array([0, 1, 0, 3])    # 预测

手动计算一下mIoU

类别0：= 交集 / 并集 = 1 / 2 = 0.5

类别1：= 1 / 1 = 1

类别2：= 0 / 1 = 0

类别3：= 1 / 1 = 1

因此 ：mIou = (0.5+1+0+1)/4 = 0.625

利用tensorflow的 API计算如下：

import numpy as np
import tensorflow as tf


a = np.array([0, 1, 2, 3])
b = np.array([0, 1, 0, 3])

iou = tf.keras.metrics.MeanIoU(num_classes=4)
iou.update_state(a, b)
print('mIoU为： ',iou.result().numpy())

上述是所有类别都参与运算的情况，现在假设我们要忽略的类别为3

以下分两种情况讨论：

情况1：

a = np.array([0, 1, 2, 3])    # 真实
b = np.array([0, 1, 0, 3])    # 预测

情况2：

a = np.array([0, 1, 2, 3])    # 真实
b = np.array([0, 1, 0, 0])    # 预测

情况1和情况2的不同点在于，真实值中类别3的位置对应的预测值为3或者其他值时，应该怎么计算单个类别的iou呢？（例如情况2中的类别0）

带着问题我们首先去用tensorflow的API来计算，然后再分析手动计算的方法

先看代码：

情况1：

a = np.array([0, 1, 2, 3])
b = np.array([0, 1, 0, 3])

iou = tf.keras.metrics.MeanIoU(num_classes=4)
sample_weight = np.ones_like(a)
sample_weight[a == 3] = 0
iou.update_state(a, b,sample_weight)
print('mIoU为： ',iou.result().numpy())

 情况2：

a = np.array([0, 1, 2, 3])
b = np.array([0, 1, 0, 0])

iou = tf.keras.metrics.MeanIoU(num_classes=4)
sample_weight = np.ones_like(a)
sample_weight[a == 3] = 0
iou.update_state(a, b,sample_weight)
print('mIoU为： ',iou.result().numpy())

 可见两种情况的mIoU计算结果不受类别3位置的元素类别影响，也就是说，在预测结果中真实标签位置的预测值随意，不影响最后的结果

因此手动计算方法可以归纳如下：

把真实标签中被忽略类别的位置记录下来，然后把真实值和预测值在该位置的元素直接删除，然后按照计算常规mIou的计算方法计算即可

下面分析一下API的使用方法：

我们要忽略某一类别，实际是将此类别计算IoU的权重设置为0。tensorflow在MeanIoU类中的update_state()方法有一个sample_weight参数：

该参数的官方解释含义如下：（贴一下连接：tf.keras.metrics.MeanIoU  |  TensorFlow Core v2.3.0）

 sample_weight是一个张量或数组，形状与y_true和y_pred相同，或者广播后相同，为了避免广播造成的错误，建议使用与y_true和y_pred相同的张量。默认情况下sample_weight即使不设置，在计算时tensorflow也会自动补充为1。当我们想忽略某一类别时，直接将sample_weight权重数组对应位置设置为0就行，其他位置的值还保持为1

这样再回头看上述代码，应该就能比较清楚了。

2、高维张量：

一维张量理解后多维的多维的也就能顺理成章的。在此只给出例子：

读者可以自己手动计算一下。

a = np.array([[[0, 1, 2], [0, 1, 2], [0, 1, 3]],[[0, 1, 1], [0, 1, 0], [0, 1, 2]]])
b = np.array([[[0, 1, 2], [0, 1, 2], [0, 1, 3]],[[0, 1, 1], [0, 0, 0], [0, 1, 2]]])
a = np.expand_dims(a,-1)
b = np.expand_dims(b,-1)


iou = tf.keras.metrics.MeanIoU(num_classes=4)
sample_weight = np.ones_like(a)
sample_weight[a == 3] = 0
iou.update_state(a, b,sample_weight)
print(a.shape)
print(b.shape)
print('mIoU为： ',iou.result().numpy())

补充：

当我们想要自定义可以忽略某一类别的mIou计算类的时候可以参考如下的自定义方法：

import cv2
import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf

a = np.array([[[0, 1, 2], [0, 1, 2], [0, 1, 3]], [[0, 1, 1], [0, 1, 0], [0, 1, 2]]])
b = np.array([[[0, 1, 2], [0, 1, 2], [0, 1, 3]], [[0, 1, 1], [0, 0, 0], [0, 1, 2]]])
a = np.expand_dims(a, -1)
b = np.expand_dims(b, -1)

iou = tf.keras.metrics.MeanIoU(num_classes=4)
sample_weight = np.ones_like(a)
sample_weight[a == 3] = 0
iou.update_state(a, b, sample_weight)
print('a的shape：', a.shape)
print('b的shape：', b.shape)
print('tensorflow 内置API计算mIoU为： ', iou.result().numpy())


class MeanIoU(tf.keras.metrics.MeanIoU):
    def __init__(self, num_classes, ignore_class_id, name=None, dtype=None):
        """
        num_classes：需要计算mIoU的像素类别数 +1   ！！！一定要加1
        ignore_class_id： 需要忽略的像素类别 类型为列表
        """
        super(MeanIoU, self).__init__(num_classes, name=name, dtype=dtype)
        self.ignore_class_id = ignore_class_id

    def __call__(self, y_true, y_pred):
        if self.ignore_class_id:
            self.sample_weight = np.ones_like(y_true.numpy())
            for i in self.ignore_class_id:
                self.sample_weight[y_true.numpy() == i] = 0
        else:
            self.sample_weight = None
        return super().__call__(y_true, y_pred, sample_weight=self.sample_weight)


iou2 = MeanIoU(4, [3, ])
a = tf.convert_to_tensor(a)
b = tf.convert_to_tensor(b)
iou2(a, b)
print("自定义的mIoU类计算结果：", iou2.result().numpy())