医学图像语义分割多分类任务计算mIoU和dice的函数

多分类分割网络输出为的结果为[channel, height, width]，channel对应分类，每个channel的pixel的值为对应当前分类的置信度。
需要先对输出结果执行一次np.argmax操作，将其转化为二维Tensor，在不同channel取置信度最高的channel(class)的值（如:0、1、2等）作为该像素的值

print(arr) # (8, 512, 512)
print(np.argmax(arr, axis=0)) # (1, 512, 512)

函数输入:

变量	简介
seg	神经网络输出的预测结果，类型:ndarray，shape:[channel=1, height, width]
gt	对应样本的ground truth，类型:ndarray，shape:[channel=1, height, width]
classes	分类数量，0到classes的数值-1
background_id	背景在预测结果的Tensor中对应分类的值(用于过滤背景)

# 使用numba加速两个函数的计算速度
# pip install numpy
# pip install numba

import numpy as np
from numba import jit
# jit编译器编译执行python代码，不使用python解析器，编译执行加速计算速度
@jit(nopython=True)
def cal_mIoU(seg, gt, classes=2, background_id=-1):
    channel_iou = []
    for i in range(classes):
        if i == background_id:
            continue
        cond = i ** 2
        # 计算相交部分
        inter = len(np.where(seg * gt == cond)[0])
        union = len(np.where(seg == i)[0]) + len(np.where(gt == i)[0]) - inter
        if union == 0:
            iou = 0
        else:
            iou = inter / union
        channel_iou.append(iou)
    res = np.array(channel_iou).mean()
    return res

@jit(nopython=True)
def cal_dice(seg, gt, classes=2, background_id=-1):
    channel_dice = []
    for i in range(classes):
        if i == background_id:
            continue
        cond = i ** 2
        # 计算相交部分
        inter = len(np.where(seg * gt == cond)[0])
        total_pix = len(np.where(seg == i)[0]) + len(np.where(gt == i)[0])
        if total_pix == 0:
            dice = 0
        else:
            dice = (2 * inter) / total_pix
        channel_dice.append(dice)
    res = np.array(channel_dice).mean()
    return res