语义分割评价指标(Dice coefficient, IoU)

import numpy as np
import torch
from matplotlib import pyplot as plt
from torchnet import meter

n_classes = 3  # 分类数
confusion_matrix = meter.ConfusionMeter(3)
score = torch.Tensor([[[[1, 3], [2, 5]], [[2, 2], [3, 7]], [[0, 1], [9, 0]]]])  # torch.Size([1, 3, 2, 2])
target = torch.tensor([[[1, 0], [2, 0]]])  # torch.Size([1, 2, 2])

# 注意2D时，cross_entropy的输入是(N,C,W,H), target是(N,W,H)
loss = torch.nn.functional.cross_entropy(score, target)  # tensor(0.7366)

# confusion_matrix要求predicted和target维度相同，且num_classes>=predicted,target>=0
predicted = score.argmax(dim=1).reshape(-1)  # torch.Size([4])
target = target.reshape(-1)  # torch.Size([4])
confusion_matrix.add(predicted, target)
cm_value = confusion_matrix.value()
# 计算全局预测准确率(混淆矩阵的对角线为预测正确的个数)
mpa = np.diag(cm_value).sum() / cm_value.sum()
# 计算每个类别的准确率
cpa = np.diag(cm_value) / cm_value.sum(1)
# 计算每个类别预测与真实目标的iou
iu = np.diag(cm_value) / (cm_value.sum(1) + cm_value.sum(0) - np.diag(cm_value))

# 绘制混淆矩阵
labels = ['cat', 'dog', 'pig']  # 每种类别的标签
# 显示数据
plt.imshow(cm_value, cmap=plt.cm.Blues)
# 在图中标注数量/概率信息
thresh = cm_value.max() / 2  # 数值颜色阈值，如果数值超过这个，就颜色加深。
for x in range(n_classes):
    for y in range(n_classes):
        # 注意这里的matrix[y, x]不是matrix[x, y]
        info = int(cm_value[y, x])
        plt.text(x, y, info,
                 verticalalignment='center',
                 horizontalalignment='center',
                 color="white" if info > thresh else "black")
plt.tight_layout()  # 保证图不重叠
plt.yticks(range(n_classes), labels)
plt.xticks(range(n_classes), labels, rotation=45)  # X轴字体倾斜45°
plt.show()
plt.close()

import torch
from torchnet import meter

def dice_coeff(x: torch.Tensor, target: torch.Tensor, ignore_index: int = -100, epsilon=1e-6):
    # Average of Dice coefficient for all batches, or for a single mask
    # 计算一个batch中所有图片某个类别的dice_coefficient
    d = 0.
    batch_size = x.shape[0]
    for i in range(batch_size):
        x_i = x[i].reshape(-1)
        t_i = target[i].reshape(-1)
        if ignore_index >= 0:
            # 找出mask中不为ignore_index的区域
            roi_mask = torch.ne(t_i, ignore_index)
            x_i = x_i[roi_mask]
            t_i = t_i[roi_mask]
        inter = torch.dot(x_i, t_i)
        sets_sum = torch.sum(x_i) + torch.sum(t_i)
        if sets_sum == 0:
            sets_sum = 2 * inter

        d += (2 * inter + epsilon) / (sets_sum + epsilon)
    return d / batch_size


def multiclass_dice_coeff(x: torch.Tensor, target: torch.Tensor, ignore_index: int = -100, epsilon=1e-6):
    """Average of Dice coefficient for all classes"""
    dice = 0.
    for channel in range(x.shape[1]):
        dice += dice_coeff(x[:, channel, ...], target[:, channel, ...], ignore_index, epsilon)

    return dice / x.shape[1]


if __name__ == '__main__':
    n_classes = 3  # 分类数
    confusion_matrix = meter.ConfusionMeter(3)
    output = torch.Tensor([[[[1, 3], [2, 5]], [[2, 2], [3, 7]], [[0, 1], [9, 0]]]])  # torch.Size([1, 3, 2, 2])
    target = torch.tensor([[[1, 0], [2, 0]]])  # torch.Size([1, 2, 2])
    # [1, 3, 2, 2] -> [1, 2, 2] -> [1, 2, 2, 3] -> [1, 3, 2, 2]
    pred = torch.nn.functional.one_hot(output.argmax(dim=1), n_classes).permute(0, 3, 1, 2).float()
    # [1, 2, 2] -> [1, 2, 2, 3] -> [1, 3, 2, 2]
    dice_target = torch.nn.functional.one_hot(target, n_classes).permute(0, 3, 1, 2).float()
    dice = multiclass_dice_coeff(pred, dice_target)
    print(dice)  # tensor(0.7778)