code: def wasserstein_disagreement_map(
      prediction, ground_truth, weight_map=None, M=None
  ):
      """
      Function to calculate the pixel-wise Wasserstein distance between the
      flattened prediction and the flattened labels (ground_truth) with respect
      to the distance matrix on the label space M.
  
      :param prediction: the logits after softmax
      :param ground_truth: segmentation ground_truth
      :param M: distance matrix on the label space
      :return: the pixelwise distance map (wass_dis_map)
      """
      if weight_map is not None:
          # raise NotImplementedError
          tf.logging.warning("Weight map specified but not used.")
  
      assert M is not None, "Distance matrix is required."
      # pixel-wise Wassertein distance (W) between flat_pred_proba and flat_labels
      # wrt the distance matrix on the label space M
      num_classes = prediction.shape[1].value
      ground_truth.set_shape(prediction.shape)
      unstack_labels = tf.unstack(ground_truth, axis=-1)
      unstack_labels = tf.cast(unstack_labels, dtype=tf.float64)
      unstack_pred = tf.unstack(prediction, axis=-1)
      unstack_pred = tf.cast(unstack_pred, dtype=tf.float64)
      # print("shape of M", M.shape, "unstacked labels", unstack_labels,
      #       "unstacked pred" ,unstack_pred)
      # W is a weighting sum of all pairwise correlations (pred_ci x labels_cj)
      pairwise_correlations = []
      for i in range(num_classes):
          for j in range(num_classes):
              pairwise_correlations.append(
                  M[i, j]
                  * tf.multiply(unstack_pred[i], unstack_labels[j])
              )
      wass_dis_map = tf.add_n(pairwise_correlations)
      return wass_dis_map
  
imports: from __future__ import absolute_import, print_function, division
  import numpy
  import tensorflow
  from niftynet.engine.application_factory import LossSegmentationFactory
  from niftynet.layer.base_layer import Layer
label_desc: Use image node_id 10 for calling a global function with example usage: wasserstein_disagreement_map(prediction, ground_truth, weight_map, M) and returns: wass_dis_map