From b1a9ad5868783b50e93817424a66bee59f8371b9 Mon Sep 17 00:00:00 2001
From: Mikael Brudfors <mbrudfors@nvidia.com>
Date: Wed, 30 Nov 2022 11:29:48 +0000
Subject: [PATCH 1/8] WIP: MetricsReloaded's binary and categorical pair-wise
 measures supports MONAI

---
 MetricsReloaded/metrics/monai_wrapper.py      | 171 +++++
 MetricsReloaded/metrics/pairwise_measures.py  | 616 +++++++++---------
 MetricsReloaded/utility/utils.py              |  39 +-
 README.rst                                    |  26 +
 setup.py                                      |   4 +-
 .../test_binary_monai_metrics copy.py         |  62 ++
 .../test_categorical_monai_metrics.py         |  35 +
 test/test_metrics/test_pairwise_measures.py   |  22 +-
 8 files changed, 667 insertions(+), 308 deletions(-)
 create mode 100644 MetricsReloaded/metrics/monai_wrapper.py
 create mode 100644 test/test_metrics/test_binary_monai_metrics copy.py
 create mode 100644 test/test_metrics/test_categorical_monai_metrics.py

diff --git a/MetricsReloaded/metrics/monai_wrapper.py b/MetricsReloaded/metrics/monai_wrapper.py
new file mode 100644
index 0000000..fa1e8d3
--- /dev/null
+++ b/MetricsReloaded/metrics/monai_wrapper.py
@@ -0,0 +1,171 @@
+"""Implementation for wrapping MetricsReloaded metrics as MONAI metrics.
+
+See test/test_metrics/{test_binary_monai_metrics copy.py, test_categorical_monai_metrics.py}
+for example use cases.
+
+"""
+try:
+    from monai.metrics import CumulativeIterationMetric
+except ImportError as e:
+    raise ImportError("MONAI is not installed, please install MetricsReloaded with pip install .[monai]")
+from monai.metrics.utils import (do_metric_reduction, is_binary_tensor)
+from monai.utils import MetricReduction
+import torch
+
+from MetricsReloaded.metrics.pairwise_measures import (
+    BinaryPairwiseMeasures,
+    MultiClassPairwiseMeasures,
+)
+
+
+def torch2numpy(x):
+    return x.cpu().detach().numpy()
+
+
+def numpy2torch(x, dtype, device):
+    return torch.as_tensor(x, dtype=dtype, device=device)
+
+
+class Metric4Monai(CumulativeIterationMetric):
+    """Allows for defining MetricsReloaded metrics as a CumulativeIterationMetric in MONAI
+    """
+    def __init__(
+        self,
+        metric_name,
+        reduction=MetricReduction.MEAN,
+        get_not_nans=False,
+    ) -> None:
+        super().__init__()
+        self.metric_name = metric_name
+        self.reduction = reduction
+        self.get_not_nans = get_not_nans
+
+    def aggregate(self, reduction=None):
+        data = self.get_buffer()
+        if not isinstance(data, torch.Tensor):
+            raise ValueError("the data to aggregate must be PyTorch Tensor.")
+
+        # do metric reduction
+        f, not_nans = do_metric_reduction(data, reduction or self.reduction)
+        return (f, not_nans) if self.get_not_nans else f
+
+
+class BinaryMetric4Monai(Metric4Monai):
+    """Wraps the binary pairwise metrics of MetricsReloaded
+    """
+    def __init__(
+        self,
+        metric_name,
+        reduction=MetricReduction.MEAN,
+        get_not_nans=False,
+    ) -> None:
+        super().__init__(metric_name=metric_name, reduction=reduction, get_not_nans=get_not_nans)
+
+    def _compute_tensor(self, y_pred, y):
+        """
+        Args:
+            y_pred: Prediction with dimensions (batch, channel, *spatial), where channel=1. The values should be binarized.
+            y: Ground-truth with dimensions (batch, channel, *spatial), where channel=1. The values should be binarized.
+
+        Raises:
+            ValueError: when `y` or `y_pred` is not a binarized tensor.
+            ValueError: when `y_pred` has less than three dimensions.
+            ValueError: when second dimension ~= 1
+        """
+        # Sanity check
+        is_binary_tensor(y_pred, "y_pred")
+        is_binary_tensor(y, "y")
+        dims = y_pred.ndimension()
+        if dims < 3:
+            raise ValueError(f"y_pred should have at least 3 dimensions (batch, channel, spatial), got {dims}.")
+        if y_pred.shape[1] != 1 or y.shape[1] != 1:
+            raise ValueError(f"y_pred.shape[1]={y_pred.shape[1]} and y.shape[1]={y.shape[1]} should be one.")
+        # Tensor parameters
+        device = y_pred.device
+        dtype = y_pred.dtype
+
+        # To numpy array
+        y_pred = torch2numpy(y_pred)
+        y = torch2numpy(y)
+
+        # Create binary pairwise metric object
+        bpm = BinaryPairwiseMeasures(
+            y_pred, y, axis=tuple(range(2, dims)), smooth_dr=1e-5,
+        )
+
+        # Is requested metric available?
+        if self.metric_name not in bpm.metrics:
+            raise ValueError(
+                f"Unsupported metric: {self.metric_name}"
+            )
+
+        # Compute metric
+        metric = bpm.metrics[self.metric_name]()
+
+        # Return metric as numpy array
+        return numpy2torch(metric, dtype, device)
+
+
+class CategoricalMetric4Monai(Metric4Monai):
+    """Wraps the categorical pairwise metrics of MetricsReloaded
+    """
+    def __init__(
+        self,
+        metric_name,
+        reduction=MetricReduction.MEAN,
+        get_not_nans=False,
+    ) -> None:
+        super().__init__(metric_name=metric_name, reduction=reduction, get_not_nans=get_not_nans)
+
+    def _compute_tensor(self, y_pred, y):
+        """
+        Args:
+            y_pred: Prediction with dimensions (batch, channel, *spatial). The values should be one-hot encoded and binarized.
+            y: Ground-truth with dimensions (batch, channel, *spatial). The values should be one-hot encoded and binarized.
+
+        Raises:
+            ValueError: when `y` or `y_pred` is not a binarized tensor.
+            ValueError: when `y_pred` has less than three dimensions.
+        """
+        # Sanity check
+        is_binary_tensor(y_pred, "y_pred")
+        is_binary_tensor(y, "y")
+        dims = y_pred.ndimension()
+        if dims < 3:
+            raise ValueError(f"y_pred should have at least 3 dimensions (batch, channel, spatial), got {dims}.")
+        # Tensor parameters
+        device = y_pred.device
+        dtype = y_pred.dtype
+        num_classes = y_pred.shape[1]
+
+        # Reshape for compatible dimension
+        y_pred = y_pred.reshape(y_pred.shape[0], y_pred.shape[1], -1)
+        y_pred = y_pred.permute((0, 2, 1))
+        y = y.reshape(y.shape[0], y.shape[1], -1)
+        y = y.permute((0, 2, 1))
+        dims = y_pred.ndimension()
+
+        # To numpy array
+        y_pred = torch2numpy(y_pred)
+        y = torch2numpy(y)
+
+        # Create binary pairwise metric object
+        bpm = MultiClassPairwiseMeasures(
+            y_pred, y, axis=tuple(range(1, dims)), smooth_dr=1e-5,
+            list_values=list(range(num_classes)), is_onehot=True,
+        )
+
+        # Is requested metric available?
+        if self.metric_name not in bpm.metrics:
+            raise ValueError(
+                f"Unsupported metric: {self.metric_name}"
+            )
+
+        # Compute metric
+        metric = bpm.metrics[self.metric_name]()
+
+        # Put back channel dimension
+        metric = metric[..., None]
+
+        # Return metric as numpy array
+        return numpy2torch(metric, dtype, device)
diff --git a/MetricsReloaded/metrics/pairwise_measures.py b/MetricsReloaded/metrics/pairwise_measures.py
index ff058db..ccb4f6c 100755
--- a/MetricsReloaded/metrics/pairwise_measures.py
+++ b/MetricsReloaded/metrics/pairwise_measures.py
@@ -54,7 +54,6 @@
 import pandas as pd
 from scipy.optimize import linear_sum_assignment as lsa
 
-
 __all__ = [
     "MultiClassPairwiseMeasures",
     "BinaryPairwiseMeasures",
@@ -70,7 +69,8 @@ class MultiClassPairwiseMeasures(object):
 
     """
 
-    def __init__(self, pred, ref, list_values, measures=[], dict_args={}):
+    def __init__(self, pred, ref, list_values,
+        measures=[], dict_args={}, smooth_dr=0, axis=None, is_onehot=False):
         self.pred = np.asarray(pred, dtype=np.int32)
         self.ref = np.asarray(ref, dtype=np.int32)
         self.dict_args = dict_args
@@ -82,46 +82,101 @@ def __init__(self, pred, ref, list_values, measures=[], dict_args={}):
             "balanced_accuracy": (self.balanced_accuracy, "BAcc"),
             "expected_cost": (self.expected_cost, "EC"),
         }
+        self.n_classes = len(self.list_values)
+
+        self.metrics = {
+            "Balanced Accuracy": self.balanced_accuracy,
+            "Weighted Cohens Kappa": self.weighted_cohens_kappa,
+            "Matthews Correlation Coefficient": self.matthews_correlation_coefficient,
+            "Expected Cost": self.expected_cost,
+            "Normalised Expected Cost": self.normalised_expected_cost,
+        }
+        self.smooth_dr = smooth_dr
+        self.axis = axis
+        if self.axis == None:
+            self.axis = (0, 1)
+        self.is_onehot = is_onehot
 
-    def expected_cost(self):
-        cm = self.confusion_matrix()
-        priors = np.sum(cm, 0) / np.sum(cm)
-        # print(priors,cm)
-        numb_perc = np.sum(cm, 0)
-        rmatrix = cm / numb_perc
-        prior_matrix = np.tile(priors, [cm.shape[0], 1])
-        priorbased_weights = 1.0 / (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"]
+    def confusion_matrix(self, return_onehot=False):
+        """
+        Provides the confusion matrix Prediction in rows, Reference in columns
+
+        :return: confusion_matrix
+        """
+        if self.is_onehot:
+            one_hot_pred = self.pred
+            one_hot_ref = self.ref
         else:
-            weights = priorbased_weights
-        print(weights, prior_matrix, rmatrix)
-        ec = np.sum(prior_matrix * weights * rmatrix)
-        return ec
+            one_hot_pred = one_hot_encode(self.pred, self.n_classes)
+            one_hot_ref = one_hot_encode(self.ref, self.n_classes)
+        confusion_matrix = np.matmul(np.swapaxes(one_hot_pred, -1, -2), one_hot_ref)
+        if return_onehot:
+            return confusion_matrix, one_hot_pred, one_hot_ref
+        else:
+            return confusion_matrix
+
+    def expectation_matrix(self, one_hot_pred=None, one_hot_ref=None):
+        """
+        Determination of the expectation matrix to be used for CK derivation
+
+        :return: expectation_matrix
+        """
+        if one_hot_pred is None:
+            one_hot_pred = one_hot_encode(self.pred, self.n_classes)
+        if one_hot_ref is None:
+            one_hot_ref = one_hot_encode(self.ref, self.n_classes)
+        pred_numb = np.sum(one_hot_pred, axis=len(one_hot_pred.shape) - 2)
+        ref_numb = np.sum(one_hot_ref, axis=len(one_hot_pred.shape) - 2)
+        n = one_hot_pred.shape[-2]
+        # print(pred_numb.shape, ref_numb.shape)
+        out = np.matmul(np.expand_dims(pred_numb, -1), np.expand_dims(ref_numb, -2)) / n
+        return out
+
+    def balanced_accuracy(self):
+        """Calculation of balanced accuracy as average of correctly classified
+        by reference class across all classes
 
-    def best_naive_ec(self):
+        .. math::
+
+            BA = \dfrac{\sum_{k=1}^{K} \dfrac{TP_k}{TP_k+FN_k}}{K}
+
+        :return: balanced_accuracy
+        """
         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
+        col_sum = np.sum(cm, axis=len(cm.shape) - 2)
+        numerator = np.diagonal(cm, axis1=len(cm.shape) - 2, axis2=len(cm.shape) - 1)
+        numerator = numerator/col_sum
+        numerator = numerator.sum(-1)
+        denominator = self.n_classes
+        return numerator / denominator
 
+    def expected_cost(self, normalise=False):
+        cm = self.confusion_matrix()
+        priors = np.sum(cm, axis=len(cm.shape) - 2, keepdims=True) \
+            / np.sum(cm, axis=self.axis, keepdims=True)
+        prior_matrix = np.tile(priors, [cm.shape[-2], 1])
+        priorbased_weights = 1.0 / (cm.shape[-2] * prior_matrix)
+        for c in range(cm.shape[-2]):
+            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)
+        numb_perc = np.sum(cm, axis=len(cm.shape) - 2, keepdims=True)
+        rmatrix = cm / numb_perc
+        ec = np.sum(prior_matrix * weights * rmatrix, axis=self.axis)
+        if normalise:
+            # Normalise with best naive expected cost
+            bnec = np.sum(weights * prior_matrix, axis=len(cm.shape) - 1)
+            bnec = np.min(bnec, axis=-1)
+            return ec / bnec
+        else:
+            return ec
 
     def normalised_expected_cost(self):
-        naive_cost = self.best_naive_ec()
-        # print(naive_cost)
-        ec = self.expected_cost()
-        print(ec, naive_cost)
-        return ec / naive_cost
+        nec = self.expected_cost(normalise=True)
+        # print(nec)
+        return nec
 
     def matthews_correlation_coefficient(self):
         """
@@ -136,21 +191,30 @@ def matthews_correlation_coefficient(self):
         .. math::
             cov_k(X,Y) = \dfrac{1}{K}\sum_{k=1}^{K}cov(X_k,Y_k)
 
+        Reference
+        Jurman, Riccadonna, Furlanello, (2012). A Comparison of MCC and CEN
+        Error Measures in MultiClass Prediction
+        https://journals.plos.org/plosone/article?id=10.1371/journal.pone.0041882
+
         :return: Matthews Correlation Coefficient
         """
-        one_hot_pred = one_hot_encode(self.pred, len(self.list_values))
-        one_hot_ref = one_hot_encode(self.ref, len(self.list_values))
-        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
+        if self.is_onehot:
+            one_hot_pred = self.pred
+            one_hot_ref = self.ref
+        else:
+            one_hot_pred = one_hot_encode(self.pred, self.n_classes)
+            one_hot_ref = one_hot_encode(self.ref, self.n_classes)
+
+        p_bar = np.mean(one_hot_pred, axis=len(one_hot_pred.shape) - 2, keepdims=True)
+        r_bar = np.mean(one_hot_ref, axis=len(one_hot_pred.shape) - 2, keepdims=True)
+
+        pp_cov = 1/self.n_classes * np.sum((one_hot_pred - p_bar)**2, axis=self.axis)
+        rr_cov = 1/self.n_classes * np.sum((one_hot_ref - r_bar)**2, axis=self.axis)
+        pr_cov = 1/self.n_classes * np.sum((one_hot_pred - p_bar) * (one_hot_ref - r_bar), axis=self.axis)
+
+        mcc = pr_cov/np.sqrt(rr_cov * pp_cov)
+
+        return mcc
 
     def chance_agreement_probability(self):
         """Determines the probability of agreeing by chance given two classifications.
@@ -165,49 +229,6 @@ def chance_agreement_probability(self):
             chance += prob_pred * prob_ref
         return chance
 
-    def confusion_matrix(self):
-        """
-        Provides the confusion matrix Prediction in rows, Reference in columns
-
-        :return: confusion_matrix
-        """
-        one_hot_pred = one_hot_encode(self.pred, len(self.list_values))
-        one_hot_ref = one_hot_encode(self.ref, len(self.list_values))
-        confusion_matrix = np.matmul(one_hot_pred.T, one_hot_ref)
-        return confusion_matrix
-
-    def balanced_accuracy(self):
-        """Calculation of balanced accuracy as average of correctly classified
-        by reference class across all classes
-
-        .. math::
-
-            BA = \dfrac{\sum_{k=1}^{K} \dfrac{TP_k}{TP_k+FN_k}}{K}
-
-        :return: balanced_accuracy
-        """
-        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
-
-        :return: expectation_matrix
-        """
-        one_hot_pred = one_hot_encode(self.pred, len(self.list_values))
-        one_hot_ref = one_hot_encode(self.ref, len(self.list_values))
-        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(np.reshape(pred_numb, [-1, 1]), np.reshape(ref_numb, [1, -1]))
-            / 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))
@@ -215,17 +236,19 @@ def weighted_cohens_kappa(self):
 
         :return: weighted_cohens_kappa
         """
-        cm = self.confusion_matrix()
-        exp = self.expectation_matrix()
+        cm, one_hot_pred, one_hot_ref = self.confusion_matrix(return_onehot=True)
+        exp = self.expectation_matrix(one_hot_pred=one_hot_pred, one_hot_ref=one_hot_ref)
         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)
+            weights = np.ones([self.n_classes, self.n_classes]) \
+                - np.eye(self.n_classes)
+        if len(cm.shape) == 3:
+            # Has batch dimension
+            weights = np.tile(weights,(cm.shape[0], 1, 1))
+        numerator = np.sum(weights * cm, axis=self.axis)
+        denominator = np.sum(weights * exp, axis=self.axis)
+        # print(numerator, denominator, cm, exp)
         return 1 - numerator / denominator
 
     def to_dict_meas(self, fmt="{:.4f}"):
@@ -247,6 +270,8 @@ def __init__(
         pixdim=None,
         empty=False,
         dict_args={},
+        axis=None,
+        smooth_dr=0,
     ):
 
         self.measures_dict = {
@@ -288,6 +313,47 @@ def __init__(
         self.pixdim = pixdim
         self.dict_args = dict_args
 
+        self.metrics = {
+            "False Positives": self.fp,
+            "False Negatives": self.fn,
+            "True Positives": self.tp,
+            "True Negatives": self.tn,
+            "Youden Index": self.youden_index,
+            "Sensitivity": self.sensitivity,
+            "Specificity": self.specificity,
+            "Balanced Accuracy": self.balanced_accuracy,
+            "Accuracy": self.accuracy,
+            "False Positive Rate": self.false_positive_rate,
+            "Normalised Expected Cost": self.normalised_expected_cost,
+            "Matthews Correlation Coefficient": self.matthews_correlation_coefficient,
+            "Cohens Kappa": self.cohens_kappa,
+            "Positive Likelihood Ratio": self.positive_likelihood_ratio,
+            "Prediction Overlaps Reference": self.pred_in_ref,
+            "Positive Predictive Value": self.positive_predictive_values,
+            "Recall": self.recall,
+            "FBeta": self.fbeta,
+            "Net Benefit Treated": self.net_benefit_treated,
+            "Negative Predictive Values": self.negative_predictive_values,
+            "Dice Score": self.dice_score,
+            "False Positives Per Image": self.fppi,
+            "Intersection Over Reference": self.intersection_over_reference,
+            "Intersection Over Union": self.intersection_over_union,
+            "Volume Difference": self.vol_diff,
+            "Topology Precision": self.topology_precision,
+            "Topology Sensitivity": self.topology_sensitivity,
+            "Centreline Dice Score": self.centreline_dsc,
+            "Boundary IoU": self.boundary_iou,
+            "Normalised Surface Distance": self.normalised_surface_distance,
+            "Average Symmetric Surface Distance": self.measured_average_distance,
+            "Mean Average Surfance Distance": self.measured_masd,
+            "Hausdorff Distance": self.measured_hausdorff_distance,
+            "xTh Percentile Hausdorff Distance": self.measured_hausdorff_distance_perc,
+        }
+        self.smooth_dr = smooth_dr
+        self.axis = axis
+        if self.axis == None:
+            self.axis = tuple(range(len(pred.shape)))
+
     def __fp_map(self):
         """
         This function calculates the false positive map
@@ -351,56 +417,56 @@ def n_pos_ref(self):
         """
         Returns the number of elements in ref
         """
-        return np.sum(self.ref)
+        return np.sum(self.ref, axis=self.axis)
 
     @CacheFunctionOutput
     def n_neg_ref(self):
         """
         Returns the number of negative elements in ref
         """
-        return np.sum(1 - self.ref)
+        return np.sum(1 - self.ref, axis=self.axis)
 
     @CacheFunctionOutput
     def n_pos_pred(self):
         """
         Returns the number of positive elements in the prediction
         """
-        return np.sum(self.pred)
+        return np.sum(self.pred, axis=self.axis)
 
     @CacheFunctionOutput
     def n_neg_pred(self):
         """
         Returns the number of negative elements in the prediction
         """
-        return np.sum(1 - self.pred)
+        return np.sum(1 - self.pred, axis=self.axis)
 
     @CacheFunctionOutput
     def fp(self):
         """
         Calculates the number of FP as sum of elements in FP_map
         """
-        return np.sum(self.__fp_map())
+        return np.sum(self.__fp_map(), axis=self.axis)
 
     @CacheFunctionOutput
     def fn(self):
         """
         Calculates the number of FN as sum of elements of FN_map
         """
-        return np.sum(self.__fn_map())
+        return np.sum(self.__fn_map(), axis=self.axis)
 
     @CacheFunctionOutput
     def tp(self):
         """
         Returns the number of true positive (TP) elements
         """
-        return np.sum(self.__tp_map())
+        return np.sum(self.__tp_map(), axis=self.axis)
 
     @CacheFunctionOutput
     def tn(self):
         """
         Returns the number of True Negative (TN) elements
         """
-        return np.sum(self.__tn_map())
+        return np.sum(self.__tn_map(), axis=self.axis)
 
     @CacheFunctionOutput
     def n_intersection(self):
@@ -419,7 +485,43 @@ def n_union(self):
             U = {\vert} Pred {\vert} + {\vert} Ref {\vert} - TP
 
         """
-        return np.sum(self.__union_map())
+        return np.sum(self.__union_map(), axis=self.axis)
+
+    @CacheFunctionOutput
+    def skeleton_versions(self, return_pred=True, return_ref=True):
+        """
+        Creates the skeletonised version of both reference and prediction
+
+        :return: skeleton_ref, skeleton_pred
+        """
+        skeleton_ref = None
+        if return_ref:
+            skeleton_ref = compute_skeleton(self.ref, axes=self.axis)
+        skeleton_pred = None
+        if return_pred:
+            skeleton_pred = compute_skeleton(self.pred, axes=self.axis)
+        return skeleton_ref, skeleton_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()
+        distance_ref = ndimage.distance_transform_edt(
+            1 - border_ref, sampling=self.pixdim
+        )
+        distance_pred = ndimage.distance_transform_edt(
+            1 - border_pred, sampling=self.pixdim
+        )
+        distance_border_pred = border_ref * distance_pred
+        distance_border_ref = border_pred * distance_ref
+        return distance_border_ref, distance_border_pred, border_ref, border_pred
 
     def youden_index(self):
         """
@@ -444,10 +546,10 @@ def sensitivity(self):
 
         :return: sensitivity
         """
-        if self.n_pos_ref() == 0:
+        if self.smooth_dr == 0 and self.n_pos_ref() == 0:
             warnings.warn("reference empty, sensitivity not defined")
             return np.nan
-        return self.tp() / self.n_pos_ref()
+        return self.tp() / (self.n_pos_ref() + self.smooth_dr)
 
     def specificity(self):
         """
@@ -462,10 +564,10 @@ def specificity(self):
 
         :return: specificity
         """
-        if self.n_neg_ref() == 0:
+        if self.smooth_dr == 0 and self.n_neg_ref() == 0:
             warnings.warn("reference all positive, specificity not defined")
             return np.nan
-        return self.tn() / self.n_neg_ref()
+        return self.tn() / (self.n_neg_ref() + self.smooth_dr)
 
     def balanced_accuracy(self):
         """
@@ -515,14 +617,14 @@ 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)
         alpha = c_fp * prior_background / (c_fn * prior_foreground)
-        print(prior_background, prior_foreground, alpha)
+        # 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
+        # print(r_fn, r_fp)
+        msk = alpha >= 1
+        ecn = np.zeros_like(alpha)
+        ecn[msk] = alpha[msk] * r_fp[msk] + r_fn[msk]
+        ecn[~msk] = r_fp[~msk] + 1 / alpha[~msk] * r_fn[~msk]
         return ecn
 
     def matthews_correlation_coefficient(self):
@@ -544,23 +646,6 @@ def matthews_correlation_coefficient(self):
         )
         return numerator / np.sqrt(denominator)
 
-    def expected_matching_ck(self):
-        list_values = np.unique(self.ref)
-        p_e = 0
-        for val in list_values:
-            p_er = np.sum(
-                np.where(
-                    self.ref == val, np.ones_like(self.ref), np.zeros_like(self.ref)
-                )
-            ) / np.prod(self.ref.shape)
-            p_es = np.sum(
-                np.where(
-                    self.pred == val, np.ones_like(self.pred), np.zeros_like(self.pred)
-                )
-            ) / np.prod(self.pred.shape)
-            p_e += p_es * p_er
-        return p_e
-
     def cohens_kappa(self):
         """
         Calculates and return the Cohen's kappa score defined as
@@ -575,7 +660,26 @@ def cohens_kappa(self):
 
         :return: CK
         """
-        p_e = self.expected_matching_ck()
+        def expected_matching():
+            list_values = np.unique(self.ref)
+            p_e = 0
+            n = 1
+            for i in self.axis:  n *= self.pred.shape[i]
+            for val in list_values:
+                p_er = np.sum(
+                    np.where(
+                        self.ref == val, np.ones_like(self.ref), np.zeros_like(self.ref)
+                    ), axis=self.axis,
+                ) / n
+                p_es = np.sum(
+                    np.where(
+                        self.pred == val, np.ones_like(self.pred), np.zeros_like(self.pred)
+                    ), axis=self.axis,
+                ) / n
+                p_e += p_es * p_er
+            return p_e
+
+        p_e = expected_matching()
         p_o = self.accuracy()
         numerator = p_o - p_e
         denominator = 1 - p_e
@@ -601,11 +705,8 @@ def pred_in_ref(self):
 
         :return: 1 if true, 0 otherwise
         """
-        intersection = np.sum(self.pred * self.ref)
-        if intersection > 0:
-            return 1
-        else:
-            return 0
+        intersection = np.sum(self.pred * self.ref, self.axis)
+        return np.where(intersection > 0, 1, 0)
 
     def positive_predictive_values(self):
         """
@@ -620,14 +721,14 @@ def positive_predictive_values(self):
 
         :return: PPV
         """
-        if self.flag_empty_pred:
+        if self.smooth_dr == 0 and self.flag_empty_pred:
             if self.flag_empty_ref:
                 warnings.warn("ref and prediction empty ppv not defined")
                 return np.nan
             else:
                 warnings.warn("prediction empty, ppv not defined but set to 0")
                 return 0
-        return self.tp() / (self.tp() + self.fp())
+        return self.tp() / (self.tp() + self.fp() + self.smooth_dr)
 
     def recall(self):
         """
@@ -635,15 +736,15 @@ def recall(self):
 
         :return: Recall = Sensitivity
         """
-        if self.n_pos_ref() == 0:
+        if self.smooth_dr == 0 and self.n_pos_ref() == 0:
             warnings.warn("reference is empty, recall not defined")
             return np.nan
-        if self.n_pos_pred() == 0:
+        if self.smooth_dr == 0 and self.n_pos_pred() == 0:
             warnings.warn(
                 "prediction is empty but ref not, recall not defined but set to 0"
             )
             return 0
-        return self.tp() / (self.tp() + self.fn())
+        return self.sensitivity()
 
     def fbeta(self):
         """
@@ -669,19 +770,19 @@ def fbeta(self):
         denominator = (
             np.square(beta) * self.positive_predictive_values() + self.recall()
         )
-        print(numerator, denominator, self.fn(), self.tp(), self.fp())
-        if np.isnan(denominator):
+        # print(numerator, denominator, self.fn(), self.tp(), self.fp())
+        if self.smooth_dr == 0 and np.isnan(denominator):
             if self.fp() + self.fn() > 0:
                 return 0
             else:
                 return 1  # Potentially modify to nan
-        elif denominator == 0:
+        elif self.smooth_dr == 0 and denominator == 0:
             if self.fp() + self.fn() > 0:
                 return 0
             else:
                 return 1  # Potentially modify to nan
         else:
-            return numerator / denominator
+            return numerator / (denominator + self.smooth_dr)
 
     def net_benefit_treated(self):
         """
@@ -700,7 +801,8 @@ def net_benefit_treated(self):
             er = self.dict_args["exchange_rate"]
         else:
             er = 1
-        n = np.size(self.pred)
+        n = 1
+        for i in self.axis:  n *= self.pred.shape[i]
         tp = self.tp()
         fp = self.fp()
         nb = tp / n - fp / n * er
@@ -713,7 +815,7 @@ def negative_predictive_values(self):
 
         :return: NPV
         """
-        if self.tn() + self.fn() == 0:
+        if self.smooth_dr == 0 and self.tn() + self.fn() == 0:
             if self.n_neg_ref() == 0:
 
                 warnings.warn(
@@ -725,7 +827,7 @@ def negative_predictive_values(self):
                     "Nothing negative in pred but should be NPV not defined but set to 0"
                 )
                 return 0
-        return self.tn() / (self.fn() + self.tn())
+        return self.tn() / (self.fn() + self.tn() + self.smooth_dr)
 
     def dice_score(self):
         """
@@ -745,13 +847,9 @@ def dice_score(self):
 
     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
+        This function returns the average number of false positives per image
         """
-        sum_per_image = np.sum(
-            np.reshape(self.__fp_map(), -1, self.ref.shape[-1]), axis=0
-        )
-        return np.mean(sum_per_image)
+        return self.__fp_map().mean(axis=self.axis)
 
     def intersection_over_reference(self):
         """
@@ -760,10 +858,10 @@ def intersection_over_reference(self):
 
         :return: IoR
         """
-        if self.flag_empty_ref:
+        if self.smooth_dr == 0 and self.flag_empty_ref:
             warnings.warn("Empty reference")
             return np.nan
-        return self.n_intersection() / self.n_pos_ref()
+        return self.n_intersection() / (self.n_pos_ref() + self.smooth_dr)
 
     def intersection_over_union(self):
         """
@@ -773,10 +871,10 @@ def intersection_over_union(self):
 
         :return: IoU
         """
-        if self.flag_empty_pred and self.flag_empty_ref:
+        if self.smooth_dr == 0 and self.flag_empty_pred and self.flag_empty_ref:
             warnings.warn("Both reference and prediction are empty")
             return np.nan
-        return self.n_intersection() / self.n_union()
+        return self.n_intersection() / (self.n_union() + self.smooth_dr)
 
     def com_dist(self):
         """
@@ -807,32 +905,6 @@ def com_dist(self):
                 )
             return com_dist
 
-    def com_ref(self):
-        """
-        This function calculates the centre of mass of the reference
-        prediction
-
-        :return: Centre of mass coordinates of reference when not empty, -1 otherwise
-        """
-        if self.flag_empty_ref:
-            return -1
-        return ndimage.center_of_mass(self.ref)
-
-    def com_pred(self):
-        """
-        This functions provides the centre of mass of the predmented element
-        :returns: -1 if empty image, centre of mass of prediction otherwise
-        """
-        if self.flag_empty_pred:
-            return -1
-        else:
-            return ndimage.center_of_mass(self.pred)
-
-    def list_labels(self):
-        if self.list_labels is None:
-            return ()
-        return tuple(np.unique(self.list_labels))
-
     def vol_diff(self):
         """
         This function calculates the ratio of difference in volume between
@@ -842,17 +914,6 @@ def vol_diff(self):
         """
         return np.abs(self.n_pos_ref() - self.n_pos_pred()) / self.n_pos_ref()
 
-    @CacheFunctionOutput
-    def skeleton_versions(self):
-        """
-        Creates the skeletonised version of both reference and prediction
-
-        :return: skeleton_ref, skeleton_pred
-        """
-        skeleton_ref = compute_skeleton(self.ref)
-        skeleton_pred = compute_skeleton(self.pred)
-        return skeleton_ref, skeleton_pred
-
     def topology_precision(self):
         """
         Calculates topology precision defined as
@@ -865,10 +926,10 @@ def topology_precision(self):
 
         :return: topology_precision
         """
-        skeleton_ref, skeleton_pred = self.skeleton_versions()
-        numerator = np.sum(skeleton_pred * self.ref)
-        denominator = np.sum(skeleton_pred)
-        print("top prec ", numerator, denominator)
+        _, skeleton_pred = self.skeleton_versions(return_ref=False)
+        numerator = np.sum(skeleton_pred * self.ref, axis=self.axis)
+        denominator = np.sum(skeleton_pred, axis=self.axis)
+        # print("top prec ", numerator, denominator)
         return numerator / denominator
 
     def topology_sensitivity(self):
@@ -883,10 +944,10 @@ def topology_sensitivity(self):
 
         :return: topology_sensitivity
         """
-        skeleton_ref, skeleton_pred = self.skeleton_versions()
-        numerator = np.sum(skeleton_ref * self.pred)
-        denominator = np.sum(skeleton_ref)
-        print("top sens ", numerator, denominator, skeleton_ref, skeleton_pred)
+        skeleton_ref, _ = self.skeleton_versions(return_pred=False)
+        numerator = np.sum(skeleton_ref * self.pred, axis=self.axis)
+        denominator = np.sum(skeleton_ref, axis=self.axis)
+        # print("top sens ", numerator, denominator, skeleton_ref, skeleton_pred)
         return numerator / denominator
 
     def centreline_dsc(self):
@@ -932,43 +993,20 @@ def boundary_iou(self):
             np.zeros_like(border_ref),
         )
 
-        intersect = np.sum(lim_dbp * lim_dbr)
+        intersect = np.sum(lim_dbp * lim_dbr, axis=self.axis)
         union = np.sum(
             np.where(
                 lim_dbp + lim_dbr > 0,
                 np.ones_like(border_ref),
                 np.zeros_like(border_pred),
-            )
+            ), axis=self.axis
         )
-        print(intersect, union)
+        # print(intersect, union)
         return intersect / union
         # 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()
-        oppose_ref = 1 - self.ref
-        oppose_pred = 1 - self.pred
-        distance_ref = ndimage.distance_transform_edt(
-            1 - border_ref, sampling=self.pixdim
-        )
-        distance_pred = ndimage.distance_transform_edt(
-            1 - border_pred, sampling=self.pixdim
-        )
-        distance_border_pred = border_ref * distance_pred
-        distance_border_ref = border_pred * distance_ref
-        return distance_border_ref, distance_border_pred, border_ref, border_pred
-
     def normalised_surface_distance(self):
         """
         Calculates the normalised surface distance (NSD) between prediction and reference
@@ -987,91 +1025,72 @@ def normalised_surface_distance(self):
         reg_pred = np.where(
             dist_pred <= tau, np.ones_like(dist_pred), np.zeros_like(dist_pred)
         )
-        numerator = np.sum(border_pred * reg_ref) + np.sum(border_ref * reg_pred)
-        denominator = np.sum(border_ref) + np.sum(border_pred)
+        numerator = np.sum(border_pred * reg_ref, axis=self.axis) + np.sum(border_ref * reg_pred, axis=self.axis)
+        denominator = np.sum(border_ref, axis=self.axis) + np.sum(border_pred, axis=self.axis)
         return numerator / denominator
 
-    def measured_distance(self):
-        """
-        This functions calculates the average symmetric distance and the
-        hausdorff distance between a prediction and a reference image
-
-        :return: hausdorff distance and average symmetric distance, hausdorff distance at perc
-        and masd
-        """
-        if "hd_perc" in self.dict_args.keys():
-            perc = self.dict_args["hd_perc"]
-        else:
-            perc = 95
-        if np.sum(self.pred + self.ref) == 0:
-            return 0, 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 + 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_perc = np.max(
-            [
-                np.percentile(ref_border_dist[self.ref + self.pred > 0], q=perc),
-                np.percentile(pred_border_dist[self.ref + self.pred > 0], q=perc),
-            ]
-        )
-        print(
-            ref_border_dist[self.ref + self.pred > 0],
-            pred_border_dist[self.ref + self.pred > 0],
-        )
-        print(
-            len(ref_border_dist[self.ref + self.pred > 0]),
-            len(pred_border_dist[self.ref + self.pred > 0]),
-        )
-
-        return hausdorff_distance, average_distance, hausdorff_distance_perc, masd
-
     def measured_average_distance(self):
         """
-        This function returns only the average distance when calculating the
-        distances between prediction and reference
+        This function returns the average symmetric surface distance (ASSD) between prediction and reference
 
         :return: assd
         """
-        return self.measured_distance()[1]
+        if self.smooth_dr == 0 and np.sum(self.pred + self.ref) == 0:
+            return 0
+        ref_border_dist, pred_border_dist, ref_border, pred_border = \
+            self.border_distance()
+        assd = \
+            (np.sum(ref_border_dist, axis=self.axis) + np.sum(pred_border_dist, axis=self.axis)) \
+            / (np.sum(pred_border + ref_border, axis=self.axis) + self.smooth_dr)
+        return assd
 
     def measured_masd(self):
-        return self.measured_distance()[3]
+        """
+        This function returns the mean average surfance distance (MASD) between prediction and reference
+
+        :return: masd
+        """
+        if self.smooth_dr == 0 and (np.sum(self.pred) == 0 or np.sum(self.ref) == 0):
+            return 0
+        ref_border_dist, pred_border_dist, ref_border, pred_border = \
+            self.border_distance()
+        masd = 0.5 * (
+            np.sum(ref_border_dist, axis=self.axis) / (np.sum(pred_border, axis=self.axis) + self.smooth_dr)
+            + np.sum(pred_border_dist, axis=self.axis) / (np.sum(ref_border, axis=self.axis) + self.smooth_dr)
+        )
+        return masd
 
     def measured_hausdorff_distance(self):
         """
-        This function returns only the hausdorff distance when calculated the
-        distances between prediction and reference
+        This function returns the Hausdorff distance between prediction and reference
 
-        :return: hausdorff_distance
+        :return: hd
         """
-        return self.measured_distance()[0]
+        ref_border_dist, pred_border_dist, _, _ = \
+            self.border_distance()
+        hd = np.max(np.maximum(ref_border_dist, pred_border_dist), axis=self.axis)
+        return hd
 
     def measured_hausdorff_distance_perc(self):
         """
-        This function returns the xth percentile hausdorff distance
+        This function returns the xth percentile Hausdorff distance between prediction and reference
 
-        :return: hausdorff_distance_perc
+        :return: hdp
         """
-        return self.measured_distance()[2]
+        if "hd_perc" in self.dict_args.keys():
+            perc = self.dict_args["hd_perc"]
+        else:
+            perc = 95
+        ref_border_dist, pred_border_dist, _, _ = \
+            self.border_distance()
+        msk = self.ref + self.pred > 0
+        ref_border_dist[~msk] = np.nan
+        pred_border_dist[~msk] = np.nan
+        hdp = np.maximum(
+            np.nanpercentile(ref_border_dist, q=perc, axis=self.axis),
+            np.nanpercentile(pred_border_dist, q=perc, axis=self.axis),
+        )
+        return hdp
 
     def to_dict_meas(self, fmt="{:.4f}"):
         result_dict = {}
@@ -1083,6 +1102,11 @@ def to_dict_meas(self, fmt="{:.4f}"):
             result_dict[key] = fmt.format(result)
         return result_dict  # trim the last comma
 
+    def list_labels(self):
+        if self.list_labels is None:
+            return ()
+        return tuple(np.unique(self.list_labels))
+
     def to_string_count(self, fmt="{:.4f}"):
         result_str = ""
         for key in self.measures_count:
diff --git a/MetricsReloaded/utility/utils.py b/MetricsReloaded/utility/utils.py
index 3eb652e..7d29cfb 100644
--- a/MetricsReloaded/utility/utils.py
+++ b/MetricsReloaded/utility/utils.py
@@ -214,12 +214,43 @@ def median_heuristic(matrix_proba):
 
 
 
-def compute_skeleton(img):
+def compute_skeleton(img, axes=None):
     """
     Computes skeleton using skimage.morphology.skeletonize
-    """
-    return skeletonize(img)
 
+    Supported input dimensions are:
+        img=(batch, channel, *spatial)
+        img=(batch, *spatial)
+        img=(*spatial)
+    where spatial can be a 2D or 3D image.
+
+    Note that, if given, this function iterates over batch and
+    channel dimensions, which is very inefficient.
+    """
+    ndimensions = len(img.shape)
+    if ndimensions < 2 or ndimensions > 5:
+        raise ValueError(
+            f"Image dimensions should be in range [2, 5], got {ndimensions}"
+        )
+    # Do we have to deal with batch/channel dimensions?
+    if axes is None:  axes = [0]
+    iter_ix = [i for i in range(0, axes[0])]
+    if not iter_ix:
+        # No
+        return skeletonize(img)
+    else:
+        # Yes
+        skeleton = np.zeros_like(img)
+        if len(iter_ix) == 1:
+            # Data has channel or batch dimension
+            for i in img.shape[0]:
+                skeleton[i, ...] = skeletonize(img[i, ...])
+        else:
+            # Data has both channel and batch dimension
+            for i0 in range(img.shape[0]):
+                for i1 in range(img.shape[1]):
+                    skeleton[i0, i1, ...] = skeletonize(img[i0, i1, ...])
+        return skeleton
 
 
 def compute_center_of_mass(img):
@@ -228,6 +259,8 @@ def compute_center_of_mass(img):
 
     :param: img as multidimensional array
     """
+    if img.sum() == 0:
+        return -1
     return ndimage.center_of_mass(img)
 
 
diff --git a/README.rst b/README.rst
index 9fbdca7..0af2fc1 100644
--- a/README.rst
+++ b/README.rst
@@ -51,6 +51,32 @@ You can alternatively install the package in editable mode:
 
 This is useful if you are developing MetricsReloaded and want to see changes in the code automatically applied to the installed library.
 
+With MONAI support
+---------
+
+Install the package as:
+
+    python -m pip install .[monai]
+
+to ensure that the MONAI dependency is installed.
+
+The MetricsReloaded metrics can then be used in, e.g., a MONAI training script as::
+
+    from MetricsReloaded.metrics.monai_wrapper import (
+        BinaryMetric4Monai,
+        CategoricalMetric4Monai,
+    )
+
+    # Example of using one of the binary pair-wise metrics
+    metric = CategoricalMetric4Monai(metric_name="Cohens Kappa")
+    metric(y_pred=y_pred, y=y)
+    value = metric.aggregate().item()
+
+    # Example of using one of the cateogrical pair-wise metrics
+    metric = CategoricalMetric4Monai(metric_name="Matthews Correlation Coefficient")
+    metric(y_pred=y_pred, y=y)
+    value = metric.aggregate().item()
+
 Overview
 ========
 
diff --git a/setup.py b/setup.py
index 211fb0e..13452f3 100755
--- a/setup.py
+++ b/setup.py
@@ -54,7 +54,9 @@
     },
     python_requires=">=3.7",
     install_requires=requirements,
-    extras_require={},
+    extras_require={
+        'monai': ['monai'],
+    },
     setup_requires=[
         "pytest-runner",
     ],
diff --git a/test/test_metrics/test_binary_monai_metrics copy.py b/test/test_metrics/test_binary_monai_metrics copy.py
new file mode 100644
index 0000000..f2d1948
--- /dev/null
+++ b/test/test_metrics/test_binary_monai_metrics copy.py	
@@ -0,0 +1,62 @@
+from MetricsReloaded.metrics.monai_wrapper import BinaryMetric4Monai
+from monai.utils import set_determinism
+import torch
+
+set_determinism(seed=0)
+
+n = 2
+c = 1
+x = 32
+y = 32
+z = 32
+
+y_pred = (torch.rand(n, c, x, y, z) > 0.5).float()
+y = (torch.rand(n, c, x, y, z) > 0.5).float()
+
+metric_names = (
+    "False Positives",
+    "False Negatives",
+    "True Positives",
+    "True Negatives",
+    "Youden Index",
+    "Sensitivity",
+    "Specificity",
+    "Balanced Accuracy",
+    "Accuracy",
+    "False Positive Rate",
+    "Normalised Expected Cost",
+    "Matthews Correlation Coefficient",
+    "Cohens Kappa",
+    "Positive Likelihood Ratio",
+    "Prediction Overlaps Reference",
+    "Positive Predictive Value",
+    "Recall",
+    "FBeta",
+    "Net Benefit Treated",
+    "Negative Predictive Values",
+    "Dice Score",
+    "False Positives Per Image",
+    "Intersection Over Reference",
+    "Intersection Over Union",
+    "Volume Difference",
+    "Topology Precision",
+    "Topology Sensitivity",
+    "Centreline Dice Score",
+    "Boundary IoU",
+    "Normalised Surface Distance",
+    "Average Symmetric Surface Distance",
+    "Mean Average Surfance Distance",
+    "Hausdorff Distance",
+    "xTh Percentile Hausdorff Distance",
+)
+
+for metric_name in metric_names:
+    print("=" * 32)
+    print(metric_name)
+    metric = BinaryMetric4Monai(metric_name=metric_name)
+    value = metric(y_pred=y_pred, y=y)
+    print(value)
+    metric(y_pred=y_pred, y=y)
+    value = metric.aggregate().item()
+    print(value)
+    metric.reset()
\ No newline at end of file
diff --git a/test/test_metrics/test_categorical_monai_metrics.py b/test/test_metrics/test_categorical_monai_metrics.py
new file mode 100644
index 0000000..fec47c1
--- /dev/null
+++ b/test/test_metrics/test_categorical_monai_metrics.py
@@ -0,0 +1,35 @@
+from MetricsReloaded.metrics.monai_wrapper import CategoricalMetric4Monai
+from monai.utils import set_determinism
+from monai.networks.utils import one_hot
+import torch
+
+set_determinism(seed=0)
+
+n = 2
+c = 1
+x = 32
+y = 32
+z = 32
+num_classes = 3
+
+y_pred = one_hot((num_classes*torch.rand(n, c, x, y, z)).round(), num_classes=num_classes + 1)
+y = one_hot((num_classes*torch.rand(n, c, x, y, z)).round(), num_classes=num_classes + 1)
+
+metric_names = (
+    "Balanced Accuracy",
+    "Weighted Cohens Kappa",
+    "Matthews Correlation Coefficient",
+    "Expected Cost",
+    "Normalised Expected Cost",
+)
+
+for metric_name in metric_names:
+    print("=" * 32)
+    print(metric_name)
+    metric = CategoricalMetric4Monai(metric_name=metric_name)
+    value = metric(y_pred=y_pred, y=y)
+    print(value)
+    metric(y_pred=y_pred, y=y)
+    value = metric.aggregate().item()
+    print(value)
+    metric.reset()
\ No newline at end of file
diff --git a/test/test_metrics/test_pairwise_measures.py b/test/test_metrics/test_pairwise_measures.py
index 3da7c3c..cb80e1a 100644
--- a/test/test_metrics/test_pairwise_measures.py
+++ b/test/test_metrics/test_pairwise_measures.py
@@ -5,7 +5,7 @@
     MultiLabelLocSegPairwiseMeasure as MLIS,
 )
 import numpy as np
-from MetricsReloaded.utility.utils import one_hot_encode
+from MetricsReloaded.utility.utils import (one_hot_encode, compute_center_of_mass)
 from numpy.testing import assert_allclose
 from sklearn.metrics import cohen_kappa_score as cks
 from sklearn.metrics import matthews_corrcoef as mcc
@@ -388,8 +388,15 @@ def test_distance_empty():
     pred = np.zeros([14, 14])
     ref = np.zeros([14, 14])
     bpm = PM(pred, ref)
-    value_test = bpm.measured_distance()
-    expected_dist = (0, 0, 0, 0)
+    expected_dist = 0
+    value_test = bpm.measured_average_distance()
+    assert_allclose(value_test, expected_dist)
+    value_test = bpm.measured_masd()
+    assert_allclose(value_test, expected_dist)
+    value_test = bpm.measured_hausdorff_distance()
+    assert_allclose(value_test, expected_dist)
+    value_test = bpm.measured_hausdorff_distance_perc()
+    expected_dist = np.nan
     assert_allclose(value_test, expected_dist)
 
 
@@ -588,14 +595,13 @@ def test_com_empty():
     ref0 = np.zeros([14, 14])
     bpm = PM(pred0, ref0)
     value_test = bpm.com_dist()
-    value_pred = bpm.com_pred()
-    value_ref = bpm.com_ref()
+    value_pred = compute_center_of_mass(pred0)
+    value_ref = compute_center_of_mass(ref0)
     expected_empty = -1
     assert value_test == expected_empty
     assert value_ref == expected_empty
     assert value_pred == expected_empty
 
-
 def test_com_dist():
     pred = np.zeros([14, 14])
     ref = np.zeros([14, 14])
@@ -603,8 +609,8 @@ def test_com_dist():
     ref[0:5, 0:5] = 1
     bpm = PM(pred, ref)
     value_dist = bpm.com_dist()
-    value_pred = bpm.com_pred()
-    value_ref = bpm.com_ref()
+    value_pred = compute_center_of_mass(pred)
+    value_ref = compute_center_of_mass(ref)
     expected_dist = 0
     expected_com = (2, 2)
     assert_allclose(value_pred, expected_com)

From 0afbb6905e1d6b17fb460ebe0bf83dd487d8d7f6 Mon Sep 17 00:00:00 2001
From: Mikael Brudfors <brudfors@gmail.com>
Date: Wed, 30 Nov 2022 15:08:32 +0000
Subject: [PATCH 2/8] Update README.rst

---
 README.rst | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/README.rst b/README.rst
index 0af2fc1..4f8880e 100644
--- a/README.rst
+++ b/README.rst
@@ -67,13 +67,15 @@ The MetricsReloaded metrics can then be used in, e.g., a MONAI training script a
         CategoricalMetric4Monai,
     )
 
-    # Example of using one of the binary pair-wise metrics
-    metric = CategoricalMetric4Monai(metric_name="Cohens Kappa")
+    # Use binary pair-wise metrics
+    metric_name = "Cohens Kappa"
+    metric = BinaryMetric4Monai(metric_name=metric_name)
     metric(y_pred=y_pred, y=y)
     value = metric.aggregate().item()
 
-    # Example of using one of the cateogrical pair-wise metrics
-    metric = CategoricalMetric4Monai(metric_name="Matthews Correlation Coefficient")
+    # Use categorical pair-wise metric
+    metric_name = "Matthews Correlation Coefficient"
+    metric = CategoricalMetric4Monai(metric_name=metric_name)
     metric(y_pred=y_pred, y=y)
     value = metric.aggregate().item()
 

From 521745ed8ddbb1a7d4c871a8723aef15da66c6f1 Mon Sep 17 00:00:00 2001
From: Mikael Brudfors <mbrudfors@nvidia.com>
Date: Thu, 8 Dec 2022 18:51:19 +0000
Subject: [PATCH 3/8] Removed MONAI wrapper (will be part of MONAI core)

---
 MetricsReloaded/metrics/monai_wrapper.py | 171 -----------------------
 1 file changed, 171 deletions(-)
 delete mode 100644 MetricsReloaded/metrics/monai_wrapper.py

diff --git a/MetricsReloaded/metrics/monai_wrapper.py b/MetricsReloaded/metrics/monai_wrapper.py
deleted file mode 100644
index fa1e8d3..0000000
--- a/MetricsReloaded/metrics/monai_wrapper.py
+++ /dev/null
@@ -1,171 +0,0 @@
-"""Implementation for wrapping MetricsReloaded metrics as MONAI metrics.
-
-See test/test_metrics/{test_binary_monai_metrics copy.py, test_categorical_monai_metrics.py}
-for example use cases.
-
-"""
-try:
-    from monai.metrics import CumulativeIterationMetric
-except ImportError as e:
-    raise ImportError("MONAI is not installed, please install MetricsReloaded with pip install .[monai]")
-from monai.metrics.utils import (do_metric_reduction, is_binary_tensor)
-from monai.utils import MetricReduction
-import torch
-
-from MetricsReloaded.metrics.pairwise_measures import (
-    BinaryPairwiseMeasures,
-    MultiClassPairwiseMeasures,
-)
-
-
-def torch2numpy(x):
-    return x.cpu().detach().numpy()
-
-
-def numpy2torch(x, dtype, device):
-    return torch.as_tensor(x, dtype=dtype, device=device)
-
-
-class Metric4Monai(CumulativeIterationMetric):
-    """Allows for defining MetricsReloaded metrics as a CumulativeIterationMetric in MONAI
-    """
-    def __init__(
-        self,
-        metric_name,
-        reduction=MetricReduction.MEAN,
-        get_not_nans=False,
-    ) -> None:
-        super().__init__()
-        self.metric_name = metric_name
-        self.reduction = reduction
-        self.get_not_nans = get_not_nans
-
-    def aggregate(self, reduction=None):
-        data = self.get_buffer()
-        if not isinstance(data, torch.Tensor):
-            raise ValueError("the data to aggregate must be PyTorch Tensor.")
-
-        # do metric reduction
-        f, not_nans = do_metric_reduction(data, reduction or self.reduction)
-        return (f, not_nans) if self.get_not_nans else f
-
-
-class BinaryMetric4Monai(Metric4Monai):
-    """Wraps the binary pairwise metrics of MetricsReloaded
-    """
-    def __init__(
-        self,
-        metric_name,
-        reduction=MetricReduction.MEAN,
-        get_not_nans=False,
-    ) -> None:
-        super().__init__(metric_name=metric_name, reduction=reduction, get_not_nans=get_not_nans)
-
-    def _compute_tensor(self, y_pred, y):
-        """
-        Args:
-            y_pred: Prediction with dimensions (batch, channel, *spatial), where channel=1. The values should be binarized.
-            y: Ground-truth with dimensions (batch, channel, *spatial), where channel=1. The values should be binarized.
-
-        Raises:
-            ValueError: when `y` or `y_pred` is not a binarized tensor.
-            ValueError: when `y_pred` has less than three dimensions.
-            ValueError: when second dimension ~= 1
-        """
-        # Sanity check
-        is_binary_tensor(y_pred, "y_pred")
-        is_binary_tensor(y, "y")
-        dims = y_pred.ndimension()
-        if dims < 3:
-            raise ValueError(f"y_pred should have at least 3 dimensions (batch, channel, spatial), got {dims}.")
-        if y_pred.shape[1] != 1 or y.shape[1] != 1:
-            raise ValueError(f"y_pred.shape[1]={y_pred.shape[1]} and y.shape[1]={y.shape[1]} should be one.")
-        # Tensor parameters
-        device = y_pred.device
-        dtype = y_pred.dtype
-
-        # To numpy array
-        y_pred = torch2numpy(y_pred)
-        y = torch2numpy(y)
-
-        # Create binary pairwise metric object
-        bpm = BinaryPairwiseMeasures(
-            y_pred, y, axis=tuple(range(2, dims)), smooth_dr=1e-5,
-        )
-
-        # Is requested metric available?
-        if self.metric_name not in bpm.metrics:
-            raise ValueError(
-                f"Unsupported metric: {self.metric_name}"
-            )
-
-        # Compute metric
-        metric = bpm.metrics[self.metric_name]()
-
-        # Return metric as numpy array
-        return numpy2torch(metric, dtype, device)
-
-
-class CategoricalMetric4Monai(Metric4Monai):
-    """Wraps the categorical pairwise metrics of MetricsReloaded
-    """
-    def __init__(
-        self,
-        metric_name,
-        reduction=MetricReduction.MEAN,
-        get_not_nans=False,
-    ) -> None:
-        super().__init__(metric_name=metric_name, reduction=reduction, get_not_nans=get_not_nans)
-
-    def _compute_tensor(self, y_pred, y):
-        """
-        Args:
-            y_pred: Prediction with dimensions (batch, channel, *spatial). The values should be one-hot encoded and binarized.
-            y: Ground-truth with dimensions (batch, channel, *spatial). The values should be one-hot encoded and binarized.
-
-        Raises:
-            ValueError: when `y` or `y_pred` is not a binarized tensor.
-            ValueError: when `y_pred` has less than three dimensions.
-        """
-        # Sanity check
-        is_binary_tensor(y_pred, "y_pred")
-        is_binary_tensor(y, "y")
-        dims = y_pred.ndimension()
-        if dims < 3:
-            raise ValueError(f"y_pred should have at least 3 dimensions (batch, channel, spatial), got {dims}.")
-        # Tensor parameters
-        device = y_pred.device
-        dtype = y_pred.dtype
-        num_classes = y_pred.shape[1]
-
-        # Reshape for compatible dimension
-        y_pred = y_pred.reshape(y_pred.shape[0], y_pred.shape[1], -1)
-        y_pred = y_pred.permute((0, 2, 1))
-        y = y.reshape(y.shape[0], y.shape[1], -1)
-        y = y.permute((0, 2, 1))
-        dims = y_pred.ndimension()
-
-        # To numpy array
-        y_pred = torch2numpy(y_pred)
-        y = torch2numpy(y)
-
-        # Create binary pairwise metric object
-        bpm = MultiClassPairwiseMeasures(
-            y_pred, y, axis=tuple(range(1, dims)), smooth_dr=1e-5,
-            list_values=list(range(num_classes)), is_onehot=True,
-        )
-
-        # Is requested metric available?
-        if self.metric_name not in bpm.metrics:
-            raise ValueError(
-                f"Unsupported metric: {self.metric_name}"
-            )
-
-        # Compute metric
-        metric = bpm.metrics[self.metric_name]()
-
-        # Put back channel dimension
-        metric = metric[..., None]
-
-        # Return metric as numpy array
-        return numpy2torch(metric, dtype, device)

From 801c1d12091bba6c78a6e8710c5711a701001e7c Mon Sep 17 00:00:00 2001
From: Mikael Brudfors <mbrudfors@nvidia.com>
Date: Thu, 8 Dec 2022 18:51:54 +0000
Subject: [PATCH 4/8] Removed optional MONAI dependency

---
 setup.py | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/setup.py b/setup.py
index 13452f3..24ea20c 100755
--- a/setup.py
+++ b/setup.py
@@ -54,9 +54,6 @@
     },
     python_requires=">=3.7",
     install_requires=requirements,
-    extras_require={
-        'monai': ['monai'],
-    },
     setup_requires=[
         "pytest-runner",
     ],

From a451bdf2ea12c8a2206c19167807e25f3898ccdd Mon Sep 17 00:00:00 2001
From: Mikael Brudfors <mbrudfors@nvidia.com>
Date: Thu, 8 Dec 2022 18:52:35 +0000
Subject: [PATCH 5/8] Removed MONA support from README

---
 README.rst | 28 ----------------------------
 1 file changed, 28 deletions(-)

diff --git a/README.rst b/README.rst
index 4f8880e..9fbdca7 100644
--- a/README.rst
+++ b/README.rst
@@ -51,34 +51,6 @@ You can alternatively install the package in editable mode:
 
 This is useful if you are developing MetricsReloaded and want to see changes in the code automatically applied to the installed library.
 
-With MONAI support
----------
-
-Install the package as:
-
-    python -m pip install .[monai]
-
-to ensure that the MONAI dependency is installed.
-
-The MetricsReloaded metrics can then be used in, e.g., a MONAI training script as::
-
-    from MetricsReloaded.metrics.monai_wrapper import (
-        BinaryMetric4Monai,
-        CategoricalMetric4Monai,
-    )
-
-    # Use binary pair-wise metrics
-    metric_name = "Cohens Kappa"
-    metric = BinaryMetric4Monai(metric_name=metric_name)
-    metric(y_pred=y_pred, y=y)
-    value = metric.aggregate().item()
-
-    # Use categorical pair-wise metric
-    metric_name = "Matthews Correlation Coefficient"
-    metric = CategoricalMetric4Monai(metric_name=metric_name)
-    metric(y_pred=y_pred, y=y)
-    value = metric.aggregate().item()
-
 Overview
 ========
 

From 5b5965b2ab555b74011e6a89c932cc69063d0008 Mon Sep 17 00:00:00 2001
From: Mikael Brudfors <mbrudfors@nvidia.com>
Date: Thu, 8 Dec 2022 18:53:41 +0000
Subject: [PATCH 6/8] Removed MONAI tests

---
 .../test_binary_monai_metrics copy.py         | 62 -------------------
 .../test_categorical_monai_metrics.py         | 35 -----------
 2 files changed, 97 deletions(-)
 delete mode 100644 test/test_metrics/test_binary_monai_metrics copy.py
 delete mode 100644 test/test_metrics/test_categorical_monai_metrics.py

diff --git a/test/test_metrics/test_binary_monai_metrics copy.py b/test/test_metrics/test_binary_monai_metrics copy.py
deleted file mode 100644
index f2d1948..0000000
--- a/test/test_metrics/test_binary_monai_metrics copy.py	
+++ /dev/null
@@ -1,62 +0,0 @@
-from MetricsReloaded.metrics.monai_wrapper import BinaryMetric4Monai
-from monai.utils import set_determinism
-import torch
-
-set_determinism(seed=0)
-
-n = 2
-c = 1
-x = 32
-y = 32
-z = 32
-
-y_pred = (torch.rand(n, c, x, y, z) > 0.5).float()
-y = (torch.rand(n, c, x, y, z) > 0.5).float()
-
-metric_names = (
-    "False Positives",
-    "False Negatives",
-    "True Positives",
-    "True Negatives",
-    "Youden Index",
-    "Sensitivity",
-    "Specificity",
-    "Balanced Accuracy",
-    "Accuracy",
-    "False Positive Rate",
-    "Normalised Expected Cost",
-    "Matthews Correlation Coefficient",
-    "Cohens Kappa",
-    "Positive Likelihood Ratio",
-    "Prediction Overlaps Reference",
-    "Positive Predictive Value",
-    "Recall",
-    "FBeta",
-    "Net Benefit Treated",
-    "Negative Predictive Values",
-    "Dice Score",
-    "False Positives Per Image",
-    "Intersection Over Reference",
-    "Intersection Over Union",
-    "Volume Difference",
-    "Topology Precision",
-    "Topology Sensitivity",
-    "Centreline Dice Score",
-    "Boundary IoU",
-    "Normalised Surface Distance",
-    "Average Symmetric Surface Distance",
-    "Mean Average Surfance Distance",
-    "Hausdorff Distance",
-    "xTh Percentile Hausdorff Distance",
-)
-
-for metric_name in metric_names:
-    print("=" * 32)
-    print(metric_name)
-    metric = BinaryMetric4Monai(metric_name=metric_name)
-    value = metric(y_pred=y_pred, y=y)
-    print(value)
-    metric(y_pred=y_pred, y=y)
-    value = metric.aggregate().item()
-    print(value)
-    metric.reset()
\ No newline at end of file
diff --git a/test/test_metrics/test_categorical_monai_metrics.py b/test/test_metrics/test_categorical_monai_metrics.py
deleted file mode 100644
index fec47c1..0000000
--- a/test/test_metrics/test_categorical_monai_metrics.py
+++ /dev/null
@@ -1,35 +0,0 @@
-from MetricsReloaded.metrics.monai_wrapper import CategoricalMetric4Monai
-from monai.utils import set_determinism
-from monai.networks.utils import one_hot
-import torch
-
-set_determinism(seed=0)
-
-n = 2
-c = 1
-x = 32
-y = 32
-z = 32
-num_classes = 3
-
-y_pred = one_hot((num_classes*torch.rand(n, c, x, y, z)).round(), num_classes=num_classes + 1)
-y = one_hot((num_classes*torch.rand(n, c, x, y, z)).round(), num_classes=num_classes + 1)
-
-metric_names = (
-    "Balanced Accuracy",
-    "Weighted Cohens Kappa",
-    "Matthews Correlation Coefficient",
-    "Expected Cost",
-    "Normalised Expected Cost",
-)
-
-for metric_name in metric_names:
-    print("=" * 32)
-    print(metric_name)
-    metric = CategoricalMetric4Monai(metric_name=metric_name)
-    value = metric(y_pred=y_pred, y=y)
-    print(value)
-    metric(y_pred=y_pred, y=y)
-    value = metric.aggregate().item()
-    print(value)
-    metric.reset()
\ No newline at end of file

From e96ba0e7be196c489aa02fb6cedbb366efacd419 Mon Sep 17 00:00:00 2001
From: Mikael Brudfors <mbrudfors@nvidia.com>
Date: Wed, 14 Dec 2022 11:51:28 +0000
Subject: [PATCH 7/8] Commented debug print statement

---
 MetricsReloaded/metrics/pairwise_measures.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/MetricsReloaded/metrics/pairwise_measures.py b/MetricsReloaded/metrics/pairwise_measures.py
index ccb4f6c..4cb7a78 100755
--- a/MetricsReloaded/metrics/pairwise_measures.py
+++ b/MetricsReloaded/metrics/pairwise_measures.py
@@ -841,8 +841,8 @@ def dice_score(self):
         elif self.dict_args["fbeta"] != 1:
             warnings.warn("Modifying fbeta option to get dice score")
             self.dict_args["fbeta"] = 1
-        else:
-            print("Already correct value for fbeta option")
+        # else:
+        #     print("Already correct value for fbeta option")
         return self.fbeta()
 
     def fppi(self):
@@ -884,14 +884,14 @@ def com_dist(self):
         :return: Euclidean distance between centre of mass when reference and prediction not empty
         -1 otherwise
         """
-        print("pred sum ", self.n_pos_pred(), "ref_sum ", self.n_pos_ref())
+        # print("pred sum ", self.n_pos_pred(), "ref_sum ", self.n_pos_ref())
         if self.flag_empty_pred or self.flag_empty_ref:
             return -1
         else:
             com_ref = compute_center_of_mass(self.ref)
             com_pred = compute_center_of_mass(self.pred)
 
-            print(com_ref, com_pred)
+            # print(com_ref, com_pred)
             if self.pixdim is not None:
                 com_dist = np.sqrt(
                     np.dot(

From 090a9a7eb33530434f02d9ab33e1bd7afbe302be Mon Sep 17 00:00:00 2001
From: Mikael Brudfors <mbrudfors@nvidia.com>
Date: Thu, 16 Feb 2023 17:22:18 +0000
Subject: [PATCH 8/8] Merge fix

---
 MetricsReloaded/metrics/pairwise_measures.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/MetricsReloaded/metrics/pairwise_measures.py b/MetricsReloaded/metrics/pairwise_measures.py
index ecef3e9..58e75d9 100755
--- a/MetricsReloaded/metrics/pairwise_measures.py
+++ b/MetricsReloaded/metrics/pairwise_measures.py
@@ -516,8 +516,8 @@ def border_distance(self):
         :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()
+        border_ref = MorphologyOps(self.ref, self.connectivity).border_map()
+        border_pred = MorphologyOps(self.pred, self.connectivity).border_map()
         distance_ref = ndimage.distance_transform_edt(
             1 - border_ref, sampling=self.pixdim
         )