|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| import numpy as np
|
|
|
| import torch
|
|
|
| def cal_dice(prediction, label, num=2):
|
| total_dice = np.zeros(num-1)
|
| for i in range(1, num):
|
| prediction_tmp = (prediction == i)
|
| label_tmp = (label == i)
|
| prediction_tmp = prediction_tmp.astype(np.float)
|
| label_tmp = label_tmp.astype(np.float)
|
|
|
| dice = 2 * np.sum(prediction_tmp * label_tmp) / (np.sum(prediction_tmp) + np.sum(label_tmp))
|
| total_dice[i - 1] += dice
|
|
|
| return total_dice
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| def dice(input, target, ignore_index=None):
|
| smooth = 1.
|
|
|
| iflat = input.clone().view(-1)
|
| tflat = target.clone().view(-1)
|
| if ignore_index is not None:
|
| mask = tflat == ignore_index
|
| tflat[mask] = 0
|
| iflat[mask] = 0
|
| intersection = (iflat * tflat).sum()
|
|
|
| return (2. * intersection + smooth) / (iflat.sum() + tflat.sum() + smooth)
|
|
|
|
|
| bce = torch.nn.BCEWithLogitsLoss(reduction='none')
|
|
|
| def _upscan(f):
|
| for i, fi in enumerate(f):
|
| if fi == np.inf: continue
|
| for j in range(1,i+1):
|
| x = fi+j*j
|
| if f[i-j] < x: break
|
| f[i-j] = x
|
|
|
| def distance_transform(bitmap):
|
| f = np.where(bitmap, 0.0, np.inf)
|
| for ibatch in range(f.shape[0]):
|
| for i in range(f.shape[1]):
|
| _upscan(f[ibatch, i, :])
|
| _upscan(f[ibatch, i,::-1])
|
| for i in range(f.shape[2]):
|
| _upscan(f[ibatch, :,i])
|
| _upscan(f[ibatch, ::-1,i])
|
| np.sqrt(f[ibatch], f[ibatch])
|
| return f
|
|
|
| def WatershedCrossEntropy(input, target):
|
|
|
|
|
| discmap = target.data.cpu()[:, 0, :, :]
|
| cupmap = target.data.cpu()[:, 1, :, :]
|
| disc_DT = distance_transform(discmap)
|
| cup_DT = distance_transform(cupmap)
|
| disc_DT = torch.from_numpy(disc_DT).float()
|
| cup_DT = torch.from_numpy(cup_DT).float()
|
|
|
| disc_DT = discmap * (1.0 - disc_DT/torch.max(disc_DT)) + 1.0
|
| cup_DT = cupmap * (1.0 - cup_DT/torch.max(cup_DT)) + 1.0
|
|
|
| disc_DT = disc_DT.cuda()
|
| cup_DT = cup_DT.cuda()
|
|
|
| CEloss = bce(input, target)
|
|
|
| return torch.mean(disc_DT* CEloss[:, 0 , :, :]+
|
| cup_DT*CEloss[:, 1 , :, :])
|
|
|
| def cross_entropy2d(input, target, weight=None, size_average=False):
|
|
|
| n, c, h, w = input.size()
|
|
|
|
|
| log_p = torch.nn.functional.log_softmax(input, dim=1)
|
|
|
| log_p = log_p.transpose(1, 2).transpose(2, 3).contiguous()
|
| log_p = log_p[target.view(n, h, w, 1).repeat(1, 1, 1, c) >= 0]
|
| log_p = log_p.view(-1, c)
|
|
|
| mask = target >= 0
|
| target = target[mask]
|
| loss = torch.nn.functional.nll_loss(log_p, target, weight=weight)
|
| loss = loss.float()
|
| if size_average:
|
| temp = mask.data.sum().float()
|
| loss = loss / temp
|
| return loss
|
|
|
|
|
| def dice_coefficient_numpy(binary_segmentation, binary_gt_label):
|
| '''
|
| Compute the Dice coefficient between two binary segmentation.
|
| Dice coefficient is defined as here: https://en.wikipedia.org/wiki/S%C3%B8rensen%E2%80%93Dice_coefficient
|
| Input:
|
| binary_segmentation: binary 2D numpy array representing the region of interest as segmented by the algorithm
|
| binary_gt_label: binary 2D numpy array representing the region of interest as provided in the database
|
| Output:
|
| dice_value: Dice coefficient between the segmentation and the ground truth
|
| '''
|
|
|
|
|
| binary_segmentation = np.asarray(binary_segmentation, dtype=bool)
|
| binary_gt_label = np.asarray(binary_gt_label, dtype=bool)
|
|
|
|
|
| intersection = np.logical_and(binary_segmentation, binary_gt_label)
|
|
|
|
|
| segmentation_pixels = float(np.sum(binary_segmentation.flatten()))
|
|
|
| gt_label_pixels = float(np.sum(binary_gt_label.flatten()))
|
|
|
| intersection = float(np.sum(intersection.flatten()))
|
|
|
| if segmentation_pixels == 0 and gt_label_pixels == 0:
|
| return 1.0
|
|
|
|
|
| dice_value = (2 * intersection + 1.0) / (1.001 + segmentation_pixels + gt_label_pixels)
|
|
|
|
|
| return dice_value
|
|
|
|
|
| def dice_coeff(pred, target, ret_arr = False):
|
| """This definition generalize to real valued pred and target vector.
|
| This should be differentiable.
|
| pred: tensor with first dimension as batch
|
| target: tensor with first dimension as batch
|
| """
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| target = target.data.cpu()
|
| if len(pred.shape) == 2:
|
| return dice_coefficient_numpy(pred, target)
|
| else:
|
| all_dice = []
|
| for i in range(pred.shape[0]):
|
| dice = dice_coefficient_numpy(pred[i, ...], target[i, ...])
|
| all_dice.append(dice)
|
|
|
| if ret_arr:
|
| return [np.array(all_dice)]
|
| return [sum(all_dice) / len(all_dice)]
|
|
|
| def dice_coeff_2label(pred, target, ret_arr = False):
|
| """This definition generalize to real valued pred and target vector.
|
| This should be differentiable.
|
| pred: tensor with first dimension as batch
|
| target: tensor with first dimension as batch
|
| """
|
|
|
| target = target.data.cpu()
|
|
|
|
|
|
|
|
|
|
|
|
|
| if len(pred.shape) == 3:
|
| return dice_coefficient_numpy(pred[0, ...], target[0, ...]), dice_coefficient_numpy(pred[1, ...], target[1, ...])
|
| else:
|
| dice_cup = []
|
| dice_disc = []
|
| for i in range(pred.shape[0]):
|
| cup, disc = dice_coefficient_numpy(pred[i, 0, ...], target[i, 0, ...]), dice_coefficient_numpy(pred[i, 1, ...], target[i, 1, ...])
|
| dice_cup.append(cup)
|
| dice_disc.append(disc)
|
| if ret_arr:
|
| return [np.array(dice_cup), np.array(dice_disc)]
|
| return [sum(dice_cup) / len(dice_cup), sum(dice_disc) / len(dice_disc)]
|
|
|
| def dice_coeff_3label(pred, target, ret_arr = False):
|
| """This definition generalize to real valued pred and target vector.
|
| This should be differentiable.
|
| pred: tensor with first dimension as batch
|
| target: tensor with first dimension as batch
|
| """
|
|
|
| target = target.data.cpu()
|
| target = np.array(target)
|
|
|
|
|
|
|
|
|
|
|
|
|
| if len(pred.shape) == 2:
|
| return dice_coefficient_numpy((pred==1).astype(float), (target==1).astype(float)), dice_coefficient_numpy((pred==2).astype(float), (target==2).astype(float)), dice_coefficient_numpy((pred==3).astype(float), (target==3).astype(float))
|
| else:
|
| dice_lv = []
|
| dice_myo = []
|
| dice_rv = []
|
| for i in range(pred.shape[0]):
|
| lv, myo, rv = dice_coefficient_numpy((pred[i]==1).astype(float), (target[i]==1).astype(float)), dice_coefficient_numpy((pred[i]==2).astype(float), (target[i]==2).astype(float)), dice_coefficient_numpy((pred[i]==3).astype(float), (target[i]==3).astype(float))
|
| dice_lv.append(lv)
|
| dice_myo.append(myo)
|
| dice_rv.append(rv)
|
| if ret_arr:
|
| return [np.array(dice_lv), np.array(dice_myo), np.array(dice_rv)]
|
| return [sum(dice_lv) / len(dice_lv), sum(dice_myo) / len(dice_myo), sum(dice_rv) / len(dice_rv)]
|
|
|
| def dice_loss(pred, target):
|
| """This definition generalize to real valued pred and target vector.
|
| This should be differentiable.
|
| pred: tensor with first dimension as batch
|
| target: tensor with first dimension as batch
|
| """
|
| return 1 - dice_coeff(pred, target)
|
|
|
|
|
| def DiceLoss(input, target):
|
| '''
|
| in tensor fomate
|
| :param input:
|
| :param target:
|
| :return:
|
| '''
|
| smooth = 1.
|
| iflat = input.contiguous().view(-1)
|
| tflat = target.contiguous().view(-1)
|
| intersection = (iflat * tflat).sum()
|
|
|
| return 1 - ((2. * intersection + smooth) /
|
| (iflat.sum() + tflat.sum() + smooth))
|
|
|
| def Balanced_DiceLoss(input, target):
|
| '''
|
|
|
| :param input:
|
| :param target:
|
| :return:
|
| '''
|
| input = torch.sigmoid(input)
|
|
|
| return 0.5 * (DiceLoss(input[:, 0, ...], target[:, 0, ...]) + DiceLoss(input[:, 1, ...], target[:, 1, ...])) |
