metrics.py

import torch
import numpy as np
from hausdorff import hausdorff_distance
from medpy.metric.binary import hd, dc

def dice(pred, target):
    pred = pred.contiguous()
    target = target.contiguous()
    smooth = 0.00001

    # intersection = (pred * target).sum(dim=2).sum(dim=2)
    pred_flat = pred.view(1, -1)
    target_flat = target.view(1, -1)

    intersection = (pred_flat * target_flat).sum().item()

    # loss = (1 - ((2. * intersection + smooth) / (pred.sum(dim=2).sum(dim=2) + target.sum(dim=2).sum(dim=2) + smooth)))
    dice = (2 * intersection + smooth) / (pred_flat.sum().item() + target_flat.sum().item() + smooth)
    return dice

def dice3D(img_gt, img_pred, voxel_size):
    """
    Function to compute the metrics between two segmentation maps given as input.

    Parameters
    ----------
    img_gt: np.array
    Array of the ground truth segmentation map.

    img_pred: np.array
    Array of the predicted segmentation map.

    voxel_size: list, tuple or np.array
    The size of a voxel of the images used to compute the volumes.

    Return
    ------
    A list of metrics in this order, [Dice LV, Volume LV, Err LV(ml),
    Dice RV, Volume RV, Err RV(ml), Dice MYO, Volume MYO, Err MYO(ml)]
    """

    if img_gt.ndim != img_pred.ndim:
        raise ValueError("The arrays 'img_gt' and 'img_pred' should have the "
                         "same dimension, {} against {}".format(img_gt.ndim,
                                                                img_pred.ndim))

    res = []
    # Loop on each classes of the input images
    for c in [3, 1, 2]:
        # Copy the gt image to not alterate the input
        gt_c_i = np.copy(img_gt)
        gt_c_i[gt_c_i != c] = 0

        # Copy the pred image to not alterate the input
        pred_c_i = np.copy(img_pred)
        pred_c_i[pred_c_i != c] = 0

        # Clip the value to compute the volumes
        gt_c_i = np.clip(gt_c_i, 0, 1)
        pred_c_i = np.clip(pred_c_i, 0, 1)

        # Compute the Dice
        dice = dc(gt_c_i, pred_c_i)

        # Compute volume
        # volpred = pred_c_i.sum() * np.prod(voxel_size) / 1000.
        # volgt = gt_c_i.sum() * np.prod(voxel_size) / 1000.

        # res += [dice, volpred, volpred-volgt]
        res += [dice]

    return res

def hd_3D(img_pred, img_gt, labels=[3, 1, 2]):
    res = []
    for c in labels:
        gt_c_i = np.copy(img_gt)
        gt_c_i[gt_c_i != c] = 0

        pred_c_i = np.copy(img_pred)
        pred_c_i[pred_c_i != c] = 0

        gt_c_i = np.clip(gt_c_i, 0, 1)
        pred_c_i = np.clip(pred_c_i, 0, 1)

        if np.sum(pred_c_i) == 0 or np.sum(gt_c_i) == 0:
            hausdorff = 0
        else:
            hausdorff = hd(pred_c_i, gt_c_i)

        res += [hausdorff]

    return res

def cal_hausdorff_distance(pred,target):

    pred = np.array(pred.contiguous())
    target = np.array(target.contiguous())
    result = hausdorff_distance(pred,target,distance="euclidean")

    return result

def make_one_hot(input, num_classes):
    """Convert class index tensor to one hot encoding tensor.
    Args:
         input: A tensor of shape [N, 1, *]
         num_classes: An int of number of class
    Returns:
        A tensor of shape [N, num_classes, *]
    """
    shape = np.array(input.shape)
    shape[1] = num_classes
    shape = tuple(shape)
    result = torch.zeros(shape).scatter_(1, input.cpu().long(), 1)
    # result = result.scatter_(1, input.cpu(), 1)

    return result

def match_pred_gt(pred, gt):
    """ pred: (1, C, H, W)
        gt: (1, C, H, W)
    """
    gt_labels = torch.unique(gt, sorted=True)[1:]
    pred_labels = torch.unique(pred, sorted=True)[1:]

    if len(gt_labels) != 0 and len(pred_labels) != 0:
        dice_Matrix = torch.zeros((len(pred_labels), len(gt_labels)))
        for i, pl in enumerate(pred_labels):
            pred_i = torch.tensor(pred==pl, dtype=torch.float)
            for j, gl in enumerate(gt_labels):
                dice_Matrix[i, j] = dice(make_one_hot(pred_i, 2)[0], make_one_hot(gt==gl, 2)[0])

        # max_axis0 = np.max(dice_Matrix, axis=0)
        max_arg0 = np.argmax(dice_Matrix, axis=0)
    else:
        return torch.zeros_like(pred)

    pred_match = torch.zeros_like(pred)
    for i, arg in enumerate(max_arg0):
        pred_match[pred==pred_labels[arg]] = i + 1
    return pred_match

if __name__ == "__main__":
    npy_path = "/home/fcheng/Cardia/source_code/logs/logs_df_50000/eval_pp_test/200.npy"
    pred_df, gt_df = np.load(npy_p)