add high level function for rigid registration

examples/registration/rigid_registration.py

@@ -1,28 +1,13 @@
-import sys, os
+import os
 import numpy as np
-import pylab as py
-
-from scipy.optimize import minimize
-
-pymirc_path = os.path.join('..','..')
-if not pymirc_path in sys.path: sys.path.append(pymirc_path)
+import matplotlib.pyplot as py
 
 import pymirc.viewer           as pymv
 import pymirc.fileio           as pymf
 import pymirc.image_operations as pymi
-import pymirc.metrics.cost_functions as pymr
-
-#---------------------------------------------------------------------------------------------
 
 data_dir = os.path.join('..','..','data','nema_petct')
 
-if not os.path.exists(data_dir):
-  url = 'https://kuleuven.box.com/s/wub9pk0yvt8kjqyj7p0bz11boca4334x'
-  print('please first download example PET/CT data from:')
-  print(url)
-  print('and unzip into: ', data_dir)
-  sys.exit()
-
 # read PET/CT nema phantom dicom data sets
 # the image data in those two data sets are aligned but
 # on different grids
@@ -38,50 +23,13 @@
 R          = pymi.kul_aff(rp, origin = np.array(ct_vol.shape)/2)
 ct_vol_rot = pymi.aff_transform(ct_vol, R, ct_vol.shape, cval = ct_vol.min()) 
 
-# define the affine transformation that maps the PET on the CT grid
-# we need this to avoid the additional interpolation from the PET to the CT grid
-pre_affine = np.linalg.inv(pet_dcm.affine) @ ct_dcm.affine
-
-#---------------------------------------------------------------------------
-#---------------------------------------------------------------------------
-# rigidly align the shifted PET to CT by minimizing neg mutual information
-# downsample the arrays by a factor for a fast pre-registration
-reg_params = np.zeros(6)
-
-# (1) initial registration with downsampled arrays
-# define the down sampling factor
-dsf    = 3
-ds_aff = np.diag([dsf,dsf,dsf,1.])
-
-ct_vol_rot_ds = pymi.aff_transform(ct_vol_rot, ds_aff, np.ceil(np.array(ct_vol.shape)/dsf).astype(int))
-
-res = minimize(pymr.regis_cost_func, reg_params, 
-               args = (ct_vol_rot_ds, pet_vol, True, True, pymr.neg_mutual_information, pre_affine @ ds_aff), 
-               method = 'Powell', 
-               options = {'ftol':1e-2, 'xtol':1e-2, 'disp':True, 'maxiter':20, 'maxfev':5000})
-
-reg_params = res.x.copy()
-# we have to scale the translations by the down sample factor since they are in voxels
-reg_params[:3] *= dsf
-
-# (2) registration with full arrays
-res = minimize(pymr.regis_cost_func, reg_params, 
-               args = (ct_vol_rot, pet_vol, True, True, pymr.neg_mutual_information, pre_affine), 
-               method = 'Powell', 
-               options = {'ftol':1e-2, 'xtol':1e-2, 'disp':True, 'maxiter':20, 'maxfev':5000})
-reg_params = res.x.copy()
-
-#---------------------------------------------------------------------------
-#---------------------------------------------------------------------------
-
-# define the final affine transformation that maps from the PET grid to the rotated CT grid
-af  = pre_affine @ pymi.kul_aff(reg_params, origin = np.array(ct_vol_rot.shape)/2)
-pet_coreg = pymi.aff_transform(pet_vol, af, ct_vol_rot.shape, cval = pet_vol.min())
+pet_coreg, coreg_aff, coreg_params = pymi.rigid_registration(pet_vol, ct_vol_rot, 
+                                                             pet_dcm.affine, ct_dcm.affine)
 
 imshow_kwargs = [{'cmap':py.cm.Greys_r, 'vmin': -200, 'vmax' : 200},
                  {'cmap':py.cm.Greys_r, 'vmin': -200, 'vmax' : 200},
                  {'cmap':py.cm.Greys, 'vmin':0, 'vmax':np.percentile(pet_coreg,99.9)}]
 
 vi = pymv.ThreeAxisViewer([ct_vol_rot, ct_vol_rot, pet_coreg], ovols = [None, pet_coreg, None],
                           imshow_kwargs = imshow_kwargs, 
-                           voxsize = ct_dcm.voxsize, width = 6)
+                          voxsize = ct_dcm.voxsize, width = 6)

pymirc/image_operations/__init__.py

@@ -9,3 +9,4 @@
 from .resample_cont_cont          import resample_cont_cont
 from .grad                        import grad, div, complex_grad, complex_div
 from .reorient                    import reorient_image_and_affine, flip_image_and_affine
+from .rigid_registration          import rigid_registration

pymirc/image_operations/rigid_registration.py

@@ -0,0 +1,96 @@
+import numpy as np
+
+from scipy.optimize import minimize
+
+import pymirc.image_operations as pymi
+import pymirc.metrics.cost_functions as pymr
+
+def rigid_registration(vol_float, vol_fixed, aff_float, aff_fixed,
+                       downsample_facs = [4], metric = pymr.neg_mutual_information,
+                       opts = {'ftol':1e-2,'xtol':1e-2,'disp':True,'maxiter':20,'maxfev':5000},
+                       method = 'Powell', metric_kwargs = {}):
+  """ rigidly coregister a 3D floating volume to a fixed (reference) volume
+
+  Parameters
+  ----------
+  vol_float : 3D numpy array
+    the floating volume
+
+  vol_fixed : 3D numpy array
+    the fixed (reference) volume
+
+  aff_float : 2D 4x4 numpy array
+    affine transformation matrix that maps from pixel to world coordinates for floating volume
+
+  aff_fixed : 2D 4x4 numpy array
+    affine transformation matrix that maps from pixel to world coordinates for fixed volume
+
+  downsample_facs : None of array_like
+    perform registrations on downsampled grids before registering original volumes
+    (multi-resolution approach)
+
+  metric : function(x,y, **kwargs)
+    metric that compares transformed floating and fixed volume
+ 
+  metric_kwargs : dict
+    keyword arguments passed to the metric function
+
+  opts : dictionary
+    passed to scipy.optimize.minimize as options
+
+  method : string
+    passed to scipy.optimize.minimize as method (which optimizer to use)
+
+  Returns
+  -------
+  tuple of 3
+  - the transformed (coregistered) floating volume
+  - the registration affine transformation matrix 
+  - the 6 registration parameters (3 translations, 3 rotations) from which the
+    affine matrix was derived
+
+  Note
+  ----
+
+  To apply the registration affine transformation use:
+  new_vol = pymi.aff_transform(vol, reg_aff, vol_fixed.shape, cval = vol.min())
+  
+  """
+  # define the affine transformation that maps the floating to the fixed voxel grid
+  # we need this to avoid the additional interpolation in case the voxel sizes are not the same
+  pre_affine = np.linalg.inv(aff_float) @ aff_fixed
+
+  reg_params = np.zeros(6)
+
+  # (1) initial registration with downsampled arrays
+  if downsample_facs is not None:
+    for dsf in downsample_facs:
+      ds_aff = np.diag([dsf,dsf,dsf,1.])
+
+      # down sample fixed volume
+      vol_fixed_ds = pymi.aff_transform(vol_fixed, ds_aff, 
+                                        np.ceil(np.array(vol_fixed.shape)/dsf).astype(int))
+
+      res = minimize(pymr.regis_cost_func, reg_params, method = method, options = opts,
+                     args = (vol_fixed_ds, vol_float, True, True, metric, pre_affine @ ds_aff,
+                             metric_kwargs))
+
+      reg_params = res.x.copy()
+      # we have to scale the translations by the down sample factor since they are in voxels
+      reg_params[:3] *= dsf
+
+
+  # (2) registration with full arrays
+  res = minimize(pymr.regis_cost_func, reg_params, method = method, options = opts, 
+                 args = (vol_fixed, vol_float, True, True, metric, pre_affine,
+                         metric_kwargs))
+  reg_params = res.x.copy()
+
+
+  # define the final affine transformation that maps from the PET grid to the rotated CT grid
+  reg_aff  = pre_affine @ pymi.kul_aff(reg_params, origin = np.array(vol_fixed.shape)/2)
+
+  # transform the floating volume
+  vol_float_coreg = pymi.aff_transform(vol_float, reg_aff, vol_fixed.shape, cval = vol_float.min())
+
+  return vol_float_coreg, reg_aff, reg_params

pymirc/metrics/cost_functions.py

@@ -131,7 +131,8 @@ def neg_mutual_information(x, y, nbins = 40, norm = True):
 
 #----------------------------------------------------------------
 def regis_cost_func(params, img_fix, img_float, verbose = False,
-                    rotate = True, metric = neg_mutual_information, pre_affine = None, metric_kwargs = {}):
+                    rotate = True, metric = neg_mutual_information, pre_affine = None, 
+                    metric_kwargs = {}):
   """Generic cost function for rigid registration
 
   Parameters