pairwise_measures.py

from __future__ import absolute_import, print_function
import warnings

import numpy as np

from utils import (
    CacheFunctionOutput, 
    MorphologyOps,
    compute_center_of_mass,
    compute_skeleton,
    one_hot_encode,
    distance_transform_edt,
)


class MultiClassPairwiseMeasures(object):
    """Class dealing with measures of direct multi-class such as MCC, Cohen's kappa or balanced accuracy"""

    def __init__(self, pred, ref, list_values, measures=[], dict_args={}):
        self.pred = np.asarray(pred, dtype=np.int32)
        self.ref = np.asarray(ref, dtype=np.int32)
        self.dict_args = dict_args
        self.list_values = list_values
        self.measures = measures
        self.measures_dict = {
            "mcc": (self.matthews_correlation_coefficient, "MCC"),
            "wck": (self.weighted_cohens_kappa, "WCK"),
            "balanced_accuracy": (self.balanced_accuracy, "BAcc"),
            "expected_cost": (self.expected_cost, "EC")
        }
        
    def expected_cost(self):
        cm = self.confusion_matrix()
        priors = np.sum(cm,0)/np.sum(cm)
        numb_perc = np.sum(cm,0)
        rmatrix = cm / numb_perc
        prior_matrix = np.tile(priors,[cm.shape[0],1])
        priorbased_weights = 1/(cm.shape[1] * prior_matrix)
        for c in range(cm.shape[0]):
            priorbased_weights[c,c] = 0
        if 'ec_costs' in self.dict_args.keys():
            weights = self.dict_args['ec_costs']
        else:
            weights = priorbased_weights
        ec = np.sum(prior_matrix * weights * rmatrix)
        return ec

    def best_naive_ec(self):
        cm = self.confusion_matrix()
        priors = np.sum(cm, 0)/np.sum(cm)
        prior_matrix = np.tile(priors,[cm.shape[0],1])
        priorbased_weights = 1/(cm.shape[1] * prior_matrix)
        for c in range(cm.shape[0]):
            priorbased_weights[c,c] = 0

        if 'ec_costs' in self.dict_args.keys():
            weights = self.dict_args['ec_costs']
        else:
            weights = priorbased_weights
        total_cost = np.sum(weights * prior_matrix,1)
        return np.min(total_cost)

    def normalised_expected_cost(self):
        naive_cost = self.best_naive_ec()
        ec = self.expected_cost()
        return ec / naive_cost

    def matthews_correlation_coefficient(self):
        one_hot_pred = self.one_hot_pred()
        one_hot_ref = self.one_hot_ref()
        cov_pred = 0
        cov_ref = 0
        cov_pr = 0
        for f in range(len(self.list_values)):
            cov_pred += np.cov(one_hot_pred[:, f], one_hot_pred[:, f])[0, 1]
            cov_ref += np.cov(one_hot_ref[:, f], one_hot_ref[:, f])[0, 1]
            cov_pr += np.cov(one_hot_pred[:, f], one_hot_ref[:, f])[0, 1]
        print(cov_pred, cov_ref, cov_pr)
        numerator = cov_pr
        denominator = np.sqrt(cov_pred * cov_ref)
        return numerator / denominator

    def chance_agreement_probability(self):
        """Determines the probability of agreeing by chance given two classifications.
        To be used for CK calculation"""
        chance = 0
        for f in self.list_values:
            prob_pred = len(np.where(self.pred == f)[0]) / np.size(self.pred)
            prob_ref = len(np.where(self.ref == f)[0]) / np.size(self.ref)
            chance += prob_pred * prob_ref
        return chance

    def confusion_matrix(self):
        """Provides the confusion matrix Prediction in rows, Reference in columns"""
        one_hot_pred = self.one_hot_pred()
        one_hot_ref = self.one_hot_ref()
        confusion_matrix = np.matmul(one_hot_pred.T, one_hot_ref)
        return confusion_matrix

    def one_hot_pred(self):
        return one_hot_encode(self.pred, np.max(self.list_values) + 1)

    def one_hot_ref(self):
        return one_hot_encode(self.ref, np.max(self.list_values) + 1)

    def balanced_accuracy(self):
        """Calculation of balanced accuracy as average of correctly classified
        by reference class across all classes"""
        cm = self.confusion_matrix()
        col_sum = np.sum(cm, 0)
        numerator = np.sum(np.diag(cm) / col_sum)
        denominator = len(self.list_values)
        return numerator / denominator

    def expectation_matrix(self):
        """Determination of the expectation matrix to be used for CK derivation"""
        one_hot_pred = self.one_hot_pred()
        one_hot_ref = self.one_hot_ref()
        pred_numb = np.sum(one_hot_pred, 0)
        ref_numb = np.sum(one_hot_ref, 0)
        print(pred_numb.shape, ref_numb.shape)
        return np.matmul(pred_numb[..., None], ref_numb[None]) \
            / np.shape(one_hot_pred)[0]

    def weighted_cohens_kappa(self):
        """Derivation of weighted cohen's kappa. The weight matrix is set to 1-ID(len(list_values))
        - cost of 1 for each error type if no weight provided
        Ref: https://en.wikipedia.org/wiki/Cohen%27s_kappa (Weighted kappa)"""
        cm = self.confusion_matrix() / len(self.pred)
        exp = self.expectation_matrix()
        if "weights" in self.dict_args.keys():
            weights = self.dict_args["weights"]
        else:
            weights = np.ones([len(self.list_values), len(self.list_values)]) \
                - np.eye(len(self.list_values))
        numerator = np.sum(weights * cm)
        denominator = np.sum(weights * exp)
        print(numerator, denominator, cm, exp)
        return 1 - numerator / denominator




class BinaryPairwiseMeasures(object):
    def __init__(
        self,
        pred,
        ref,
        measures=[],
        num_neighbors=8,
        pixdim=None,
        empty=False,
        dict_args={},
    ):

        self.measures_dict = {
            "numb_ref": (self.n_pos_ref, "NumbRef"),
            "numb_pred": (self.n_pos_pred, "NumbPred"),
            "numb_tp": (self.n_intersection, "NumbTP"),
            "numb_fp":(self.fp, "NumbFP"),
            "numb_fn": (self.fn, "NumbFN"),
            "accuracy": (self.accuracy, "Accuracy"),
            "expected_cost": (self.normalised_expected_cost, "ECn"),
            "balanced_accuracy": (self.balanced_accuracy, "BalAcc"),
            "cohens_kappa": (self.cohens_kappa, "CohensKappa"),
            "lr+": (self.positive_likelihood_ratio, "LR+"),
            "iou": (self.intersection_over_union, "IoU"),
            "fbeta": (self.fbeta, "FBeta"),
            "youden_ind": (self.youden_index, "YoudenInd"),
            "mcc": (self.matthews_correlation_coefficient, "MCC"),
            "centreline_dsc": (self.centreline_dsc, "CentreLineDSC"),
            "assd": (self.measured_average_distance, "ASSD"),
            "boundary_iou": (self.boundary_iou, "BoundaryIoU"),
            "hd": (self.measured_hausdorff_distance, "HD"),
            "hd_perc": (self.measured_hausdorff_distance_perc, "HDPerc"),
            "masd": (self.measured_masd, "MASD"),
            "nsd": (self.normalised_surface_distance, "NSD"),
        }
        
        self.pred = np.where(pred, 1, 0)  # binarise
        self.ref = np.where(ref, 1, 0)  # binarise
        self.flag_empty = empty
        self.measures = measures if measures is not None else self.measures_dict
        self.neigh = num_neighbors
        self.pixdim = pixdim
        self.dict_args = dict_args

    def __fp_map(self):
        """
        This function calculates the false positive map
        :return: FP map
        """
        return np.asarray((self.pred - self.ref) > 0.0, dtype=np.float32)

    def __fn_map(self):
        """
        This function calculates the false negative map
        :return: FN map
        """
        return np.asarray((self.ref - self.pred) > 0.0, dtype=np.float32)

    def __tp_map(self):
        """
        This function calculates the true positive map
        :return: TP map
        """
        return np.asarray((self.ref + self.pred) > 1.0, dtype=np.float32)

    def __tn_map(self):
        """
        This function calculates the true negative map
        :return: TN map
        """
        return np.asarray((self.ref + self.pred) < 0.5, dtype=np.float32)

    def __union_map(self):
        """
        This function calculates the union map between predicted and
        reference image
        :return: union map
        """
        return np.asarray((self.ref + self.pred) > 0.5, dtype=np.float32)

    def __intersection_map(self):
        """
        This function calculates the intersection between predicted and
        reference image
        :return: intersection map
        """
        return np.multiply(self.ref, self.pred)

    @CacheFunctionOutput
    def n_pos_ref(self):
        return np.sum(self.ref)

    @CacheFunctionOutput
    def n_neg_ref(self):
        return np.sum(1 - self.ref)

    @CacheFunctionOutput
    def n_pos_pred(self):
        return np.sum(self.pred)

    @CacheFunctionOutput
    def n_neg_pred(self):
        return np.sum(1 - self.pred)

    @CacheFunctionOutput
    def fp(self):
        return np.sum(self.__fp_map())

    @CacheFunctionOutput
    def fn(self):
        return np.sum(self.__fn_map())

    @CacheFunctionOutput
    def tp(self):
        return np.sum(self.__tp_map())

    @CacheFunctionOutput
    def tn(self):
        return np.sum(self.__tn_map())

    @CacheFunctionOutput
    def n_intersection(self):
        return np.sum(self.__intersection_map())

    @CacheFunctionOutput
    def n_union(self):
        return np.sum(self.__union_map())

    def youden_index(self):
        return self.sensitivity() + self.specificity() - 1

    def sensitivity(self):
        if self.n_pos_ref() == 0:
            warnings.warn("reference empty, sensitivity not defined")
            return np.nan
        return self.tp() / self.n_pos_ref()

    def precision(self):
        if self.n_pos_pred() == 0:
            if self.n_pos_ref() == 0:
                warnings.warn("ref and prediction empty precision not defined")
                return np.nan
            else:
                warnings.warn("prediction empty, precision not defined but set to 0")
                return 0
        return self.tp() / (self.tp() + self.fp())

    def specificity(self):
        if self.n_neg_ref() == 0:
            warnings.warn("reference all positive, specificity not defined")
            return np.nan
        return self.tn() / self.n_neg_ref()

    def balanced_accuracy(self):
        return 0.5 * self.sensitivity() + 0.5 * self.specificity()

    def accuracy(self):
        return (self.tn() + self.tp()) / (self.tn() + self.tp() + self.fn() + self.fp())

    def false_positive_rate(self):
        return self.fp() / self.n_neg_ref()

    def normalised_expected_cost(self):
        prior_background = (self.tn() + self.fp())/np.size(self.ref)
        prior_foreground = (self.tp() + self.fn())/np.size(self.ref)

        if 'cost_fn' in self.dict_args.keys():
            c_fn = self.dict_args['cost_fn']
        else:
            c_fn = 1.0/(2*prior_foreground)
        if 'cost_fp' in self.dict_args.keys():
            c_fp = self.dict_args['cost_fp']
        else:
            c_fp = 1.0/(2*prior_background)
        
        alpha = c_fp * prior_background / (c_fn * prior_foreground)
        print(prior_background, prior_foreground, alpha)
        r_fp = self.fp()/self.n_neg_ref()
        r_fn = self.fn()/self.n_pos_ref()
        print(r_fn, r_fp)
        if alpha >= 1:
            ecn = alpha * r_fp + r_fn
        else:
            ecn = r_fp + 1/alpha * r_fn
        return ecn

    def matthews_correlation_coefficient(self):
        numerator = self.tp() * self.tn() - self.fp() * self.fn()
        denominator = (
            (self.tp() + self.fp())
            * (self.tp() + self.fn())
            * (self.tn() + self.fp())
            * (self.tn() + self.fn())
        )
        return numerator / np.sqrt(denominator)

    def expected_matching_ck(self):
        # https://en.wikipedia.org/wiki/Cohen%27s_kappa#Definition (just two raters)        
        list_values = np.unique(self.ref)
        p_e = 0
        for val in list_values:
            p_er = np.sum(self.ref == val) / np.prod(self.ref.shape)
            p_es = np.sum(self.pred == val) / np.prod(self.pred.shape)
            p_e += p_es * p_er
        return p_e

    def cohens_kappa(self):
        p_e = self.expected_matching_ck()
        p_o = self.accuracy()
        numerator = p_o - p_e
        denominator = 1 - p_e
        return numerator / denominator

    def positive_likelihood_ratio(self):
        numerator = self.sensitivity()
        denominator = 1 - self.specificity()
        return numerator / denominator

    def pred_in_ref(self):
        intersection = np.sum(self.pred * self.ref)
        if intersection > 0:
            return 1
        else:
            return 0
        
    def fbeta(self):
        if "beta" in self.dict_args.keys():
            beta = self.dict_args["beta"]
        else:
            beta = 1
        numerator = (
            (1 + np.square(beta)) * self.precision() * self.sensitivity()
        )
        denominator = (
            np.square(beta) * self.precision() + self.sensitivity()
        )
        print(numerator, denominator, self.fn(), self.tp(), self.fp())
        if np.isnan(denominator):
            if self.fp() + self.fn() > 0:
                return 0
            else:
                return 1  # Potentially modify to nan
        elif denominator == 0:
            if self.fp() + self.fn() > 0:
                return 0
            else:
                return 1  # Potentially modify to nan
        else:
            return numerator / denominator

    def negative_predictive_values(self):
        """
        This function calculates the negative predictive value ratio between
        the number of true negatives and the total number of negative elements
        :return:
        """
        if self.tn() + self.fn() == 0:
            if self.n_neg_ref() == 0:

                warnings.warn(
                    "Nothing negative in either pred or ref, NPV not defined and set to nan"
                )
                return np.nan  # Potentially modify to 1
            else:
                warnings.warn(
                    "Nothing negative in pred but should be NPV not defined but set to 0"
                )
                return 0
        return self.tn() / (self.fn() + self.tn())

    def dice_score(self):
        """
        This function returns the dice score coefficient between a reference
        and predicted images
        :return: dice score
        """
        return 2 * self.tp() / np.sum(self.ref + self.pred)

    def fppi(self):
        """
        This function returns the average number of false positives per
         image, assuming that the cases are collated on the last axis of the array
        """
        sum_per_image = np.sum(
            np.reshape(self.__fp_map(), (-1, self.ref.shape[-1])), axis=0
        )
        return np.mean(sum_per_image)

    def intersection_over_reference(self):
        """
        This function returns the intersection over reference ratio
        :return:
        """
        return self.n_intersection() / self.n_pos_ref()

    def intersection_over_union(self):
        """
        This function the intersection over union ratio - Definition of
        jaccard coefficient
        :return:
        """
        return self.n_intersection() / self.n_union()

    def com_dist(self):
        """
        This function calculates the euclidean distance between the centres
        of mass of the reference and prediction.
        :return:
        """
        print("pred sum ", self.n_pos_pred(), "ref_sum ", self.n_pos_ref())
        if self.flag_empty:
            return -1
        else:
            com_ref = self.center_of_mass_ref()
            com_pred = self.center_of_mass_pred()
            print(com_ref, com_pred)
            if self.pixdim is not None:
                com_dist = np.sqrt(
                    np.dot(
                        np.square(np.asarray(com_ref) - np.asarray(com_pred)),
                        np.square(self.pixdim),
                    )
                )
            else:
                com_dist = np.sqrt(
                    np.sum(np.square(np.asarray(com_ref) - np.asarray(com_pred)))
                )
            return com_dist

    def center_of_mass_ref(self):
        """
        This function calculates the centre of mass of the reference image
        :return:
        """
        return compute_center_of_mass(self.ref)

    def center_of_mass_pred(self):
        """
        This functions provides the centre of mass of the predicted image
        :return:
        """
        if self.flag_empty:
            return -1
        else:
            return compute_center_of_mass(self.pred)

    def vol_diff(self):
        """
        This function calculates the ratio of difference in volume between
        the reference and predicted images.
        :return: vol_diff
        """
        return np.abs(self.n_pos_ref() - self.n_pos_pred()) / self.n_pos_ref()

    def skeletonize_ref(self):
        return compute_skeleton(self.ref)

    def skeletonize_pred(self):
        return compute_skeleton(self.pred)

    def topology_precision(self):
        skeleton_pred = self.skeletonize_pred()
        numerator = np.sum(skeleton_pred * self.ref)
        denominator = np.sum(skeleton_pred)
        print("top prec ", numerator, denominator)
        return numerator / denominator

    def topology_sensitivity(self):
        skeleton_ref = self.skeletonize_ref()
        numerator = np.sum(skeleton_ref * self.pred)
        denominator = np.sum(skeleton_ref)
        print("top sens ", numerator, denominator)
        return numerator / denominator

    def centreline_dsc(self):
        top_prec = self.topology_precision()
        top_sens = self.topology_sensitivity()
        numerator = 2 * top_sens * top_prec
        denominator = top_sens + top_prec
        return numerator / denominator

    def boundary_iou(self):
        """
        This functions determines the boundary iou
        """
        border_ref = MorphologyOps(self.ref, self.neigh).border_map()
        border_pred = MorphologyOps(self.pred, self.neigh).border_map()
        return np.sum(border_ref * border_pred) / (
            np.sum(border_ref) + np.sum(border_pred)
        )

    @CacheFunctionOutput
    def border_distance(self):
        """
        This functions determines the map of distance from the borders of the
        prediction and the reference and the border maps themselves
        :return: distance_border_ref, distance_border_pred, border_ref,
        border_pred
        """
        border_ref = MorphologyOps(self.ref, self.neigh).border_map()
        border_pred = MorphologyOps(self.pred, self.neigh).border_map()
        # Compute Euclidean distance transform
        distance_ref = distance_transform_edt(1 - border_ref, sampling=self.pixdim)
        distance_pred = distance_transform_edt(1 - border_pred, sampling=self.pixdim)
        distance_pred = border_ref * distance_pred
        distance_ref = border_pred * distance_ref
        return distance_ref, distance_pred, border_ref, border_pred

    def normalised_surface_distance(self):
        if "nsd" in self.dict_args.keys():
            tau = self.dict_args["nsd"]
        else:
            tau = 1
        dist_ref, dist_pred, border_ref, border_pred = self.border_distance()
        reg_ref = np.where(dist_ref <= tau, 1, 0)
        reg_pred = np.where(dist_pred <= tau, 1, 0)
        numerator = np.sum(border_pred * reg_ref) + np.sum(border_ref * reg_pred)
        denominator = np.sum(border_ref) + np.sum(border_pred)
        return numerator / denominator

    def measured_distance(self, perc=95):
        """
        This functions calculates the average symmetric distance and the
        hausdorff distance between a predmentation and a reference image
        :return: hausdorff distance and average symmetric distance
        """
        if np.sum(self.pred + self.ref) == 0:
            return 0, 0, 0
        ref_border_dist, pred_border_dist, ref_border, pred_border \
            = self.border_distance()
        average_distance = \
            (np.sum(ref_border_dist) + np.sum(pred_border_dist)) \
            / (np.sum(pred_border) + np.sum(ref_border))
        masd = 0.5 * (
            np.sum(ref_border_dist) / np.sum(pred_border)
            + np.sum(pred_border_dist) / np.sum(ref_border)
        )
        print(
            np.sum(ref_border_dist) / np.sum(pred_border),
            np.sum(pred_border_dist) / np.sum(ref_border),
            np.sum(pred_border),
            np.sum(ref_border),
            np.sum(pred_border_dist),
            np.sum(ref_border_dist),
        )
        hausdorff_distance = np.max([np.max(ref_border_dist), np.max(pred_border_dist)])
        hausdorff_distance_95 = np.max(
            [
                np.percentile(ref_border_dist[ref_border_dist > 0], q=perc),
                np.percentile(pred_border_dist[pred_border_dist > 0], q=perc),
            ]
        )
        return hausdorff_distance, average_distance, hausdorff_distance_95, masd

    def measured_average_distance(self):
        """
        This function returns only the average distance when calculating the
        distances between predmentation and reference
        :return:
        """
        return self.measured_distance()[1]

    def measured_masd(self):
        """
        Mean average surface distance (MASD)
        """
        return self.measured_distance()[3]

    def measured_hausdorff_distance(self):
        """
        This function returns only the hausdorff distance when calculated the
        distances between predmentation and reference
        :return:
        """
        return self.measured_distance()[0]

    def measured_hausdorff_distance_perc(self):
        if "hd" in self.dict_args.keys():
            perc = self.dict_args["hd"]
        else:
            perc = 95
        return self.measured_distance(perc)[2]