医学图像语义分割多分类任务计算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

你可能感兴趣的:(Python,深度学习,神经网络,深度学习,python,语义分割)