多分类分割网络输出为的结果为[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