text: #!/usr/bin/env python
  
  """
  Example script to test a segmentation network trained in an unsupervised fashion,
  using a probabilistic atlas and unlabeled scans.
  
  Unsupervised deep learning for Bayesian brain MRI segmentation
  A.V. Dalca, E. Yu, P. Golland, B. Fischl, M.R. Sabuncu, J.E. Iglesias
  Under Review. arXiv https://arxiv.org/abs/1904.11319
  """
  
  import os
  import argparse
  import numpy as np
  import tensorflow as tf
  from tensorflow.keras import backend as K
  import voxelmorph as vxm
  
  
  # parse the commandline
  parser = argparse.ArgumentParser()
  parser.add_argument('image', help='input image to test')
  parser.add_argument('seg', help='output segmentation file')
  parser.add_argument('--model', required=True, help='keras model file')
  parser.add_argument('--atlas', required=True, help='atlas npz file')
  parser.add_argument('--mapping', required=True, help='atlas mapping filename')
  parser.add_argument('--gpu', help='GPU number - if not supplied, CPU is used')
  parser.add_argument("--max-feats", default=21, help='number of label posteriors to compute on GPU at once')
  parser.add_argument('--warped-atlas', help='save warped atlas to output vol file')
  parser.add_argument('--posteriors', help='save posteriors to output vol file')
  parser.add_argument('--warp', help='save warp to output vol file')
  parser.add_argument('--stats', help='save stats to output npz file')
  args = parser.parse_args()
  
  # load reference atlas (group labels in tissue types if necessary)
  atlas_full = vxm.py.utils.load_volfile(args.atlas, add_batch_axis=True)
  mapping = np.load(args.mapping)['mapping'].astype('int').flatten()
  assert len(mapping) == atlas_full.shape[-1], \
      'mapping shape %d is inconsistent with atlas shape %d' % (len(mapping), atlas_full.shape[-1])
  nb_labels = int(1 + np.max(mapping))
  atlas = np.zeros([*atlas_full.shape[:-1], nb_labels])
  for i in range(np.max(mapping.shape)):
      atlas[0, ..., mapping[i]] = atlas[0, ..., mapping[i]] + atlas_full[0, ..., i]
  
  # get input shape
  inshape = atlas.shape[1:-1]
  
  # load input scan
  image, affine = vxm.py.utils.load_volfile(args.image, add_batch_axis=True, add_feat_axis=True, ret_affine=True)
  
  
  # define an isolated method of computing posteriors
  def make_k_functions(vol_shape, mapping, max_feats=None, norm_post=True):
      """
      Utility to build keras (gpu-runnable) functions that will warp the atlas and compute
      posteriors given the original full atlas and unnormalized log likelihood.
  
      norm_post (True): normalize posterior? Thi sis faster on GPU, so if possible should be set to True
      max_feats (None): since atlas van be very large, warping full atlas can run out of memory on current GPUs. 
          Providing a number here will avoid OOM error, and will return several keras functions that each provide 
          the posterior computation for at most max_feats nb_full_labels. Stacking the result of calling these
          functions will provide an *unnormalized* posterior (since it can't properly normalize)
      """
      nb_full_labels = len(mapping)
      nb_labels = np.max(mapping) + 1
  
      # compute maximum features and whether to return single function
      return_single_fn = max_feats is None
      if max_feats is None:
          max_feats = nb_full_labels
      else:
          assert not norm_post, 'cannot do normalized posterior if providing max_feats'
  
      # prepare ull and 
      ull = tf.placeholder(tf.float32, vol_shape + (nb_labels, ))
      input_flow = tf.placeholder(tf.float32, vol_shape + (len(vol_shape), ))
      ul_pred = K.exp(ull)
      
      funcs = []
      for i in range(0, nb_full_labels, max_feats):
          end = min(i + max_feats, nb_full_labels)
          this_atlas_full_shape = vol_shape + (end - i, )
          this_mapping = mapping[i:end]
  
          # prepare atlas input
          input_atlas = tf.placeholder(tf.float32, this_atlas_full_shape)
  
          # warp atlas
          warped = vxm.tf.ne.transform(input_atlas, input_flow, interp_method='linear', indexing='ij')
  
          # normalized posterior
          post_lst = [ul_pred[..., this_mapping[j]] * warped[..., j] for j in range(len(this_mapping))]
          posterior = K.stack(post_lst, -1)
          
          funcs.append(K.function([input_atlas, ull, input_flow], [posterior, warped]))
  
      if return_single_fn:
          if norm_post:
              posterior = posterior / (1e-12 + K.sum(posterior, -1, keepdims=True))
          funcs = funcs.append(K.function([input_atlas, ull, input_flow], [posterior, warped]))
  
      return funcs
  
  
  # tensorflow device handling
  device, nb_devices = vxm.tf.utils.setup_device(args.gpu)
  
  with tf.device(device):
  
      # get k-functions for posterior
      funcs = make_k_functions(inshape, mapping, max_feats=args.max_feats, norm_post=False)
  
      # load the model from file
      model = vxm.networks.ProbAtlasSegmentation.load(args.model).get_gaussian_warp_model()
  
      # predict log likelihood and flow
      outputs = model.predict([image, atlas])
  
      # remove unused batch dimension
      ull_pred, mus, sigmas, flow = [array[0, ...] for array in outputs]
  
      # run through label groups and compute posteriors
      posteriors = []
      warped_atlas = []
      total_labels = atlas_full.shape[-1]
      for li, i in enumerate(range(0, total_labels, args.max_feats)):
          slc = slice(i, min(i + args.max_feats, total_labels))
          po, wa = funcs[li]([atlas_full[0, ..., slc], ull_pred, flow])
          posteriors.append(po)
          warped_atlas.append(wa)
      posteriors = np.concatenate(posteriors, -1)
      warped_atlas = np.concatenate(warped_atlas, -1)
  
      # compute final segmentation
      # note: this requires about 0.5 seconds on the CPU. This can also be computed on the GPU
      # as part of the funcs, (which would be very fast), but would require the full posterior
      # estimation on the GPU in a single run. Some GPUs will have difficulty with this.
      segmentation = posteriors.argmax(-1)
  
  # save segmentation
  vxm.py.utils.save_volfile(segmentation.astype('int32'), args.seg, affine)
  
  # save warped atlas
  if args.warped_atlas:
      vxm.py.utils.save_volfile(warped_atlas, args.warped_atlas, affine)
  
  # save posteriors
  if args.posteriors:
      normalized = posteriors / (1e-12 + np.sum(posteriors, -1, keepdims=True))
      vxm.py.utils.save_volfile(normalized, args.posteriors, affine)
  
  # save warp
  if args.warp:
      vxm.py.utils.save_volfile(flow, args.warp, affine)
  
  # save computed stats
  if args.stats:
      np.savez_compressed(args.stats, means=mus, log_variances=sigmas)
  