code: import nibabel as nib
  from tqdm import tqdm
  from scipy.ndimage import label, generate_binary_structure
  from pathlib import Path
  import json
  import numpy as np
  from ivadomed import postprocessing as imed_postpro
  from typing import List
  
  
  def run_uncertainty(image_folder):
      """Compute uncertainty from model prediction.
  
      This function loops across the model predictions (nifti masks) and estimates the uncertainty from the Monte Carlo
      samples. Both voxel-wise and structure-wise uncertainty are estimates.
  
      Args:
          image_folder (str): Folder containing the Monte Carlo samples.
      """
      # list subj_acq prefixes
      subj_acq_lst = [file.name.split('_pred')[0] for file in Path(image_folder).iterdir()
                      if file.name.endswith('.nii.gz') and '_pred' in file.name]
      # remove duplicates
      subj_acq_lst = list(set(subj_acq_lst))
      # keep only the images where unc has not been computed yet
      subj_acq_lst = [file for file in subj_acq_lst if not Path(image_folder, file + '_unc-cv.nii.gz').is_file()]
  
      # loop across subj_acq
      for subj_acq in tqdm(subj_acq_lst, desc="Uncertainty Computation"):
          # hard segmentation from MC samples
          fname_pred: Path = Path(image_folder, subj_acq + '_pred.nii.gz')
          # fname for soft segmentation from MC simulations
          fname_soft: Path = Path(image_folder, subj_acq + '_soft.nii.gz')
          # find Monte Carlo simulations
          fname_pred_lst: List[str] = []
          for file in Path(image_folder).iterdir():
              if subj_acq + '_pred_' in file.name and ('_painted' not in file.name) and ('_color' not in file.name):
                  fname_pred_lst.append(str(file))
  
          # if final segmentation from Monte Carlo simulations has not been generated yet
          if not fname_pred.is_file() or not fname_soft.is_file():
              # threshold used for the hard segmentation
              thr = 1. / len(fname_pred_lst)  # 1 for all voxels where at least on MC sample predicted 1
              # average then argmax
              combine_predictions(fname_pred_lst, str(fname_pred), str(fname_soft), thr=thr)
  
          fname_unc_vox = Path(image_folder, subj_acq + '_unc-vox.nii.gz')
          if not fname_unc_vox.is_file():
              # compute voxel-wise uncertainty map
              voxelwise_uncertainty(fname_pred_lst, str(fname_unc_vox))
  
          fname_unc_struct = Path(image_folder, subj_acq + '_unc.nii.gz')
          if not Path(image_folder, subj_acq + '_unc-cv.nii.gz').is_file():
              # compute structure-wise uncertainty
              structurewise_uncertainty(fname_pred_lst, str(fname_pred), str(fname_unc_vox), str(fname_unc_struct))
  
  
  def combine_predictions(fname_lst, fname_hard, fname_prob, thr=0.5):
      """Combine predictions from Monte Carlo simulations.
  
      Combine predictions from Monte Carlo simulations and save the resulting as:
          (1) `fname_prob`, a soft segmentation obtained by averaging the Monte Carlo samples.
          (2) `fname_hard`, a hard segmentation obtained thresholding with `thr`.
  
      Args:
          fname_lst (list of str): List of the Monte Carlo samples.
          fname_hard (str): Filename for the output hard segmentation.
          fname_prob (str): Filename for the output soft segmentation.
          thr (float): Between 0 and 1. Used to threshold the soft segmentation and generate the hard segmentation.
      """
      # collect all MC simulations
      mc_data = np.array([nib.load(fname).get_fdata() for fname in fname_lst])
      first_file_header = nib.load(fname_lst[0]).header
  
      # average over all the MC simulations
      data_prob = np.mean(mc_data, axis=0)
      # save prob segmentation
      nib_prob = nib.Nifti1Image(
          dataobj=data_prob,
          affine=first_file_header.get_best_affine(),
          header=first_file_header.copy()
      )
      nib.save(nib_prob, fname_prob)
  
      # argmax operator
      data_hard = imed_postpro.threshold_predictions(data_prob, thr=thr).astype(np.uint8)
      # save hard segmentation
      nib_hard = nib.Nifti1Image(
          dataobj=data_hard,
          affine=first_file_header.get_best_affine(),
          header=first_file_header.copy()
      )
      nib.save(nib_hard, fname_hard)
  
  
  def voxelwise_uncertainty(fname_lst, fname_out, eps=1e-5):
      """Estimate voxel wise uncertainty.
  
      Voxel-wise uncertainty is estimated as entropy over all N MC probability maps, and saved in `fname_out`.
  
      Args:
          fname_lst (list of str): List of the Monte Carlo samples.
          fname_out (str): Output filename.
          eps (float): Epsilon value to deal with np.log(0).
      """
      # collect all MC simulations
      mc_data = np.array([nib.load(fname).get_fdata() for fname in fname_lst])
      affine = nib.load(fname_lst[0]).header.get_best_affine()
  
      # entropy
      unc = np.repeat(np.expand_dims(mc_data, -1), 2, -1)  # n_it, x, y, z, 2
      unc[..., 0] = 1 - unc[..., 1]
      unc = -np.sum(np.mean(unc, 0) * np.log(np.mean(unc, 0) + eps), -1)
  
      # Clip values to 0
      unc[unc < 0] = 0
  
      # save uncertainty map
      nib_unc = nib.Nifti1Image(unc, affine)
      nib.save(nib_unc, fname_out)
  
  
  def structurewise_uncertainty(fname_lst, fname_hard, fname_unc_vox, fname_out):
      """Estimate structure wise uncertainty.
  
      Structure-wise uncertainty from N MC probability maps (`fname_lst`) and saved in `fname_out` with the following
      suffixes:
  
          * '-cv.nii.gz': coefficient of variation
          * '-iou.nii.gz': intersection over union
          * '-avgUnc.nii.gz': average voxel-wise uncertainty within the structure.
  
      Args:
          fname_lst (list of str): List of the Monte Carlo samples.
          fname_hard (str): Filename of the hard segmentation, which is used to compute the `avgUnc` by providing a mask
              of the structures.
          fname_unc_vox (str): Filename of the voxel-wise uncertainty, which is used to compute the `avgUnc`.
          fname_out (str): Output filename.
      """
      # 18-connectivity
      bin_struct = np.array(generate_binary_structure(3, 2))
  
      # load hard segmentation
      nib_hard = nib.load(fname_hard)
      data_hard = nib_hard.get_fdata()
      # Label each object of each class
      data_hard_labeled = [label(data_hard[..., i_class], structure=bin_struct)[0] for i_class in
                           range(data_hard.shape[-1])]
  
      # load all MC simulations (in mc_dict["mc_data"]) and label them (in mc_dict["mc_labeled"])
      mc_dict = {"mc_data": [], "mc_labeled": []}
      for fname in fname_lst:
          data = nib.load(fname).get_fdata()
          mc_dict["mc_data"].append([data[..., i_class] for i_class in range(data.shape[-1])])
  
          labeled_list = [label(data[..., i_class], structure=bin_struct)[0] for i_class in range(data.shape[-1])]
          mc_dict["mc_labeled"].append(labeled_list)
  
      # load uncertainty map
      data_uncVox = nib.load(fname_unc_vox).get_fdata()
  
      # Init output arrays
      data_iou, data_cv, data_avgUnc = np.zeros(data_hard.shape), np.zeros(data_hard.shape), np.zeros(data_hard.shape)
  
      # Loop across classes
      for i_class in range(data_hard.shape[-1]):
          # Hard segmentation of the i_class that has been labeled
          data_hard_labeled_class = data_hard_labeled[i_class]
          # Get number of objects in
          l, l_count = np.unique(data_hard_labeled_class, return_counts=True)
  
          # Get all non zero labels and exclude structure of 1 pixel
          labels = l[l_count != 1][1:]
          # Loop across objects
          for i_obj in labels:
              # select the current structure, remaining voxels are set to zero
              data_hard_labeled_class_obj = (np.array(data_hard_labeled_class) == i_obj).astype(int)
  
              # Get object coordinates
              xx_obj, yy_obj, zz_obj = np.where(data_hard_labeled_class_obj)
  
              # Loop across the MC samples and mask the structure of interest
              data_class_obj_mc = []
              for i_mc in range(len(fname_lst)):
                  # Get index of the structure of interest in the MC sample labeled
                  i_mc_labels, i_mc_counts = np.unique(data_hard_labeled_class_obj * mc_dict["mc_labeled"][i_mc][i_class],
                                                       return_counts=True)
                  i_mc_label = i_mc_labels[np.argmax(i_mc_counts[1:]) + 1] if len(i_mc_counts) > 1 else 0
  
                  data_tmp = np.zeros(mc_dict["mc_data"][i_mc][i_class].shape)
                  # If i_mc_label is zero, it means the structure is not present in this mc_sample
                  if i_mc_label > 0:
                      data_tmp[mc_dict["mc_labeled"][i_mc][i_class] == i_mc_label] = 1.
  
                  data_class_obj_mc.append(data_tmp.astype(np.bool))
  
              # COMPUTE IoU
              # Init intersection and union
              intersection = np.logical_and(data_class_obj_mc[0], data_class_obj_mc[1])
              union = np.logical_or(data_class_obj_mc[0], data_class_obj_mc[1])
              # Loop across remaining MC samples
              for i_mc in range(2, len(data_class_obj_mc)):
                  intersection = np.logical_and(intersection, data_class_obj_mc[i_mc])
                  union = np.logical_or(union, data_class_obj_mc[i_mc])
              # Compute float
              iou = np.sum(intersection) * 1. / np.sum(union)
              # assign uncertainty value to the structure
              data_iou[xx_obj, yy_obj, zz_obj, i_class] = iou
  
              # COMPUTE COEFFICIENT OF VARIATION
              # List of volumes for each MC sample
              vol_mc_lst = [np.sum(data_class_obj_mc[i_mc]) for i_mc in range(len(data_class_obj_mc))]
              # Mean volume
              mu_mc = np.mean(vol_mc_lst)
              # STD volume
              sigma_mc = np.std(vol_mc_lst)
              # Coefficient of variation
              cv = sigma_mc / mu_mc
              # assign uncertainty value to the structure
              data_cv[xx_obj, yy_obj, zz_obj, i_class] = cv
  
              # COMPUTE AVG VOXEL WISE UNC
              avgUnc = np.mean(data_uncVox[xx_obj, yy_obj, zz_obj, i_class])
              # assign uncertainty value to the structure
              data_avgUnc[xx_obj, yy_obj, zz_obj, i_class] = avgUnc
  
      # save nifti files
      fname_iou = fname_out.split('.nii.gz')[0] + '-iou.nii.gz'
      fname_cv = fname_out.split('.nii.gz')[0] + '-cv.nii.gz'
      fname_avgUnc = fname_out.split('.nii.gz')[0] + '-avgUnc.nii.gz'
  
      nib_iou = nib.Nifti1Image(
          dataobj=data_iou,
          affine=nib_hard.header.get_best_affine(),
          header=nib_hard.header.copy()
      )
      nib_cv = nib.Nifti1Image(
          dataobj=data_cv,
          affine=nib_hard.header.get_best_affine(),
          header=nib_hard.header.copy()
      )
      nib_avgUnc = nib.Nifti1Image(
          data_avgUnc,
          affine=nib_hard.header.get_best_affine(),
          header=nib_hard.header.copy()
      )
  
      nib.save(nib_iou, fname_iou)
      nib.save(nib_cv, fname_cv)
      nib.save(nib_avgUnc, fname_avgUnc)
  
apis: numpy.logical_or numpy.sum numpy.zeros scipy.ndimage.label numpy.logical_and numpy.argmax numpy.expand_dims numpy.array numpy.std numpy.where numpy.unique numpy.mean scipy.ndimage.generate_binary_structure