dbdicom 0.2.1__py3-none-any.whl → 0.2.3__py3-none-any.whl

dbdicom/utils/vreg.py

@@ -1,11 +1,10 @@
 import time
 import numpy as np
-import scipy.stats as stats
+import scipy
 import scipy.optimize as opt
 import scipy.ndimage as ndi
-from scipy.interpolate import interpn
+from scipy.interpolate import interpn, griddata
 from scipy.spatial.transform import Rotation
-import sklearn
 
 import pyvista as pv
 
@@ -17,6 +16,34 @@ from skimage.draw import ellipsoid
 ##########################
 
 
+def fill_gaps(data, loc, mask=None):
+    # Fill gaps in data by interpolation
+    # data is an array with values
+    # loc is a mask array defining where to interpolate (0=interpolate, 1=value)
+    # mask is a mask array to restrict the interpolation to certain regions.
+
+    x, y, z = np.indices(data.shape)
+
+    # Get locations and values of non-zero pixels
+    i = loc>0.5
+    if mask is not None:
+        i = i & (mask==1)
+    points = np.column_stack([x[i], y[i], z[i]])
+    values = data[i]
+
+    # Interpolate using griddata
+    k = loc<0.5
+    if mask is not None:
+        k = k & (mask==1)
+    filled = data.copy()
+    filled[k] = griddata(points, values, (x[k], y[k], z[k]), method='linear', fill_value=0)
+
+    if mask is not None:
+        filled *= mask
+
+    return filled
+
+
 def volume_coordinates(shape, position=[0,0,0]):
 
     # data are defined at the middle of the voxels - use p+0.5 offset.
@@ -116,7 +143,6 @@ def interpolator_displacement(displacement, shape, interpolator):
     return displacement_interp
 
 
-
 def surface_coordinates(shape):
 
     # data are defined at the edge of volume - extend shape with 1.
@@ -158,7 +184,7 @@ def pv_contour(values, data, affine, surface=False):
     surf_data = np.reshape(surf_data, surf_shape)
 
     rotation, translation, pixel_spacing = affine_components(affine)
-    grid = pv.UniformGrid(dimensions=surf_shape, spacing=pixel_spacing)
+    grid = pv.ImageData(dimensions=surf_shape, spacing=pixel_spacing)
     surf = grid.contour(values, surf_data.flatten(order="F"), method='marching_cubes')
     surf = surf.rotate_vector(rotation, np.linalg.norm(rotation)*180/np.pi, inplace=False)
     surf = surf.translate(translation, inplace=False)
@@ -299,6 +325,63 @@ def envelope(d, affine):
     return output_shape, output_pos
 
 
+def bounding_box(array:np.ndarray, affine:np.ndarray, margin:float=0)->tuple:
+    """Return a bounding box around a region of interest.
+
+    For a given mask array with know affine this finds the bounding box around the mask, and returns its shape and affine.
+
+    Args:
+        array (np.ndarray): 3d binary array with mask data (1=inside, 0=outside).
+        affine (np.ndarray): affine of the 3d volume as a 4x4 numpy array.
+        margin (float): margin (in mm) to include around the bounding box. (default=0)
+
+    Returns:
+        tuple: shape and affine of the bounding box, where shape is a tuple of 3 values and affine is a 4x4 numpy array.
+    """
+
+    _, _, pixel_spacing = affine_components(affine)
+    xmargin = np.around(margin/pixel_spacing[0]).astype(np.int16)
+    ymargin = np.around(margin/pixel_spacing[1]).astype(np.int16)
+    zmargin = np.around(margin/pixel_spacing[2]).astype(np.int16)
+
+    # Find shape and location of the box in array coordinates and get array.
+    x, y, z = np.where(array != 0)
+    x0, x1 = np.amin(x)-xmargin, np.amax(x)+xmargin
+    y0, y1 = np.amin(y)-ymargin, np.amax(y)+ymargin
+    z0, z1 = np.amin(z)-zmargin, np.amax(z)+zmargin
+    x0, x1 = np.amax([0, x0]), np.amin([x1, array.shape[0]-1])
+    y0, y1 = np.amax([0, y0]), np.amin([y1, array.shape[1]-1])
+    z0, z1 = np.amax([0, z0]), np.amin([z1, array.shape[2]-1])
+    nx = 1 + np.ceil(x1-x0).astype(np.int16)
+    ny = 1 + np.ceil(y1-y0).astype(np.int16)
+    nz = 1 + np.ceil(z1-z0).astype(np.int16)
+    box_array = array[x0:x0+nx, y0:y0+ny, z0:z0+nz]
+
+    # Get the corner in absolute coordinates and offset the affine.
+    nd = 3
+    matrix = affine[:nd,:nd]
+    offset = affine[:nd, nd]
+    r0 = np.array([x0,y0,z0])
+    r0 = np.dot(r0, matrix.T) + offset
+    box_affine = affine.copy()
+    box_affine[:nd, nd] = r0
+
+    return box_array, box_affine
+
+
+def mask_volume(array, affine, array_mask, affine_mask, margin:float=0):
+
+    # Overlay the mask on the array.
+    array_mask, _ = affine_reslice(array_mask, affine_mask, affine, array.shape)
+
+    # Mask out array pixels outside of region.
+    array *= array_mask
+
+    # Extract bounding box around non-zero pixels in masked array.
+    array, affine = bounding_box(array, affine, margin)
+    
+    return array, affine
+
 
 def apply_affine(affine, coord):
     """Apply affine transformation to an array of coordinates"""
@@ -446,6 +529,7 @@ def affine_output_geometry(input_shape, input_affine, transformation):
     return output_shape, output_affine
 
 
+
 ####################################
 ## Affine transformation and reslice
 ####################################
@@ -471,7 +555,38 @@ def affine_transform(input_data, input_affine, transformation, reshape=False, **
     return output_data, output_affine
 
 
-# TODO This needs to become a private helper function
+def affine_reslice_slice_by_slice(input_data, input_affine, output_affine, output_shape=None, slice_thickness=None, mask=False, label=False, **kwargs):
+    # generalizes affine_reslice - also works with multislice volumes where slice thickness is less than slice gap
+
+    # If 3D volume - do normal affine_reslice
+    if slice_thickness is None:
+        output_data, output_affine = affine_reslice(input_data, input_affine, output_affine, output_shape=output_shape, **kwargs)
+    # If slice thickness equals slice spacing:
+    # then its a 3D volume - do normal affine_reslice 
+    elif slice_thickness == np.linalg.norm(input_affine[:3,2]):
+        output_data, output_affine = affine_reslice(input_data, input_affine, output_affine, output_shape=output_shape, **kwargs)
+    # If multislice - perform affine slice by slice
+    else:
+        output_data = None
+        for z in range(input_data.shape[2]):
+            input_data_z, input_affine_z = extract_slice(input_data, input_affine, z, slice_thickness=slice_thickness)
+            output_data_z, output_affine = affine_reslice(input_data_z, input_affine_z, output_affine, output_shape=output_shape, **kwargs)
+            if output_data is None:
+                output_data = output_data_z
+            else:
+                output_data += output_data_z
+    # If source is a mask array, convert to binary:
+    if mask:
+        output_data[output_data > 0.5] = 1
+        output_data[output_data <= 0.5] = 0
+    # If source is a label array, convert to integers:
+    elif label:
+        output_data = np.around(output_data)
+
+    return output_data, output_affine
+    
+    
+
 def affine_reslice(input_data, input_affine, output_affine, output_shape=None, **kwargs):
 
     # If 2d array, add a 3d dimension of size 1
@@ -646,60 +761,41 @@ def freeform_deformation_align(input_data, displacement, output_shape=None, outp
     return output_data
 
 
-def extract_slice(array, affine, z, slice_thickness=None):
+def affine_slice(affine, z, slice_thickness=None):
 
     # Get the slice and its affine
-    array_z = array[:,:,z]
     affine_z = affine.copy()
     affine_z[:3,3] += z*affine[:3,2]
     # Set the slice spacing to equal the slice thickness.
     # Note: both are equal for 3D array but different for 2D multislice
     if slice_thickness is not None:
         slice_spacing = np.linalg.norm(affine[:3,2])
-        affine_z[:3,2] *= slice_thickness[z]/slice_spacing
-
-    return array_z, affine_z
-
-    
-# def _transform_slice_by_slice(
-#         moving = None, 
-#         moving_affine = None, 
-#         static_shape = None, 
-#         static_affine = None, 
-#         parameters = None, 
-#         transformation = translate, 
-#         slice_thickness = None):
-    
-#     # Note this does not work for center of mass rotation because weight array has different center of mass.
-            
-#     nz = moving.shape[2]
-#     if slice_thickness is not None:
-#         if not isinstance(slice_thickness, list):
-#             slice_thickness = [slice_thickness]*nz
-
-#     weight = np.zeros(static_shape + (nz,))
-#     coregistered = np.zeros(static_shape + (nz,))
+        if np.isscalar(slice_thickness):
+            affine_z[:3,2] *= slice_thickness/slice_spacing
+        else:
+            affine_z[:3,2] *= slice_thickness[z]/slice_spacing
+    return affine_z
 
-#     for z in range(nz):
-#         moving_z, moving_affine_z = extract_slice(moving, moving_affine, z, slice_thickness)
-#         weight[...,z] = transformation(np.ones(moving.shape[:2]), moving_affine_z, static_shape, static_affine, parameters[z])
-#         coregistered[...,z] = transformation(moving_z, moving_affine_z, static_shape, static_affine, parameters[z])
 
-#     # Average each pixel value over all slices that have sampled it
-#     coregistered = np.sum(weight*coregistered, axis=-1)
-#     weight = np.sum(weight, axis=-1)
-#     nozero = np.where(weight > 0)
-#     coregistered[nozero] = coregistered[nozero]/weight[nozero]
+def extract_slice(array, affine, z, slice_thickness=None):
+    array_z = array[:,:,z]
+    affine_z = affine_slice(affine, z, slice_thickness=slice_thickness)
+    return array_z, affine_z
 
-#     return coregistered
+def inslice_translation(affine, translation):
+    row_cosine = affine[:3,0]/np.linalg.norm(affine[:3,0])
+    column_cosine = affine[:3,1]/np.linalg.norm(affine[:3,1])
+    return translation[0]*row_cosine + translation[1]*column_cosine
 
 
 
 ####################################
-# wrappers for use in align function
+# extensions for use in align function
 ####################################
 
-
+def translate_inslice(input_data, input_affine, output_shape, output_affine, translation, **kwargs):
+    transformation = affine_matrix(translation=inslice_translation(input_affine, translation))
+    return affine_transform_and_reslice(input_data, input_affine, output_shape, output_affine, transformation, **kwargs)
 
 def translate(input_data, input_affine, output_shape, output_affine, translation, **kwargs):
     transformation = affine_matrix(translation=translation)
@@ -776,7 +872,6 @@ def freeform_align(input_data, input_affine, output_shape, output_affine, parame
 def transform_slice_by_slice(input_data, input_affine, output_shape, output_affine, parameters, transformation=translate, slice_thickness=None):
     
     # Note this does not work for center of mass rotation because weight array has different center of mass.
-            
     nz = input_data.shape[2]
     if slice_thickness is not None:
         if not isinstance(slice_thickness, list):
@@ -795,93 +890,97 @@ def transform_slice_by_slice(input_data, input_affine, output_shape, output_affi
     # Average each pixel value over all slices that have sampled it
     nozero = np.where(weight > 0)
     coregistered[nozero] = coregistered[nozero]/weight[nozero]
-
     return coregistered
 
+def passive_translation_slice_by_slice(input_affine, parameters):
+    output_affine = []
+    for z, pz in enumerate(parameters):
+        input_affine_z = affine_slice(input_affine, z)
+        rigid_transform = affine_matrix(translation=pz)
+        transformed_input_affine = rigid_transform.dot(input_affine_z)
+        output_affine.append(transformed_input_affine)
+    return output_affine
+
+def passive_rigid_transform_slice_by_slice(input_affine, parameters):
+    output_affine = []
+    for z, pz in enumerate(parameters):
+        input_affine_z = affine_slice(input_affine, z)
+        rigid_transform = affine_matrix(rotation=pz[:3], translation=pz[3:])
+        #rigid_transform = affine_matrix(rotation=z*(np.pi/180)*np.array([1,0,0])) # for debugging
+        transformed_input_affine = rigid_transform.dot(input_affine_z)
+        output_affine.append(transformed_input_affine)
+    return output_affine
+
+def passive_translation(input_affine, parameters):
+    rigid_transform = affine_matrix(translation=parameters)
+    output_affine = rigid_transform.dot(input_affine)
+    return output_affine
+
+def passive_rigid_transform(input_affine, parameters):
+    rigid_transform = affine_matrix(rotation=parameters[:3], translation=parameters[3:])
+    output_affine = rigid_transform.dot(input_affine)
+    return output_affine
+
+def multislice_to_singleslice_affine(affine_ms, slice_thickness):
+    # In a multi-slice affine, the z-spacing is the slice spacing (distance between slice centers)
+    # In a single-slice affine, the z-spacing is the slice thickness
+    affine_ss = affine_ms.copy()
+    affine_ss[:3,2] *= slice_thickness/np.linalg.norm(affine_ss[:3,2])
+    return affine_ss
+
+
 
 # default metrics
 # ---------------
 
-def sum_of_squares(static, transformed, nan=None):
+def interaction(static, transformed, nan=None):
+    # static is here a mask image.
     if nan is not None:
         i = np.where(transformed != nan)
-        return np.sum(np.square(static[i]-transformed[i]))
+        msk, img = static[i], transformed[i]
     else:
-        return np.sum(np.square(static-transformed))
+        msk, img = static, transformed
+    return np.std(img[np.where(msk>0.5)])
+    return 1/np.mean(np.abs(msk*img))
+    return np.exp(-np.mean(np.abs(msk*img)))
+    i = img[np.where(msk>0.5)]
+    return np.exp(-np.mean(i**2))
 
 
-def mutual_information(static, transformed, nan=None):
+def sum_of_squares(static, transformed, nan=None):
     if nan is not None:
         i = np.where(transformed != nan)
-        return sklearn.metrics.mutual_info_score(static[i], transformed[i])
+        st, tr = static[i], transformed[i]
     else:
-        return sklearn.metrics.mutual_info_score(static, transformed)
+        st, tr = static, transformed
+    return np.sum(np.square(st-tr))
     
 
-def normalized_mutual_information(static, transformed, nan=None):
+def mutual_information(static, transformed, nan=None):
+
+    # Mask if needed
     if nan is not None:
         i = np.where(transformed != nan)
-        return sklearn.metrics.normalized_mutual_info_score(static[i], transformed[i])
-    else:
-        return sklearn.metrics.normalized_mutual_info_score(static, transformed)
-
-
-
-# Generic objective function
-def _objective_function(
-        params, 
-        moving: np.ndarray, 
-        static: np.ndarray, 
-        transformation, 
-        metric,
-        moving_affine,
-        static_affine,
-        moving_mask, 
-        static_mask, 
-        transformation_args, 
-        metric_args):
-
-    # # Get static size (TODO: use shape instead of size in transformations)
-    # _, _, static_pixel_spacing = affine_components(static_affine)
-    # static_size = np.multiply(np.array(static.shape)-1, static_pixel_spacing)
-
-    # Transform the moving image
-    nan = 0
-    if transformation_args is None:
-        transformed = transformation(moving, moving_affine, static.shape, static_affine, params, cval=nan)
-    else:
-        transformed = transformation(moving, moving_affine, static.shape, static_affine, params, transformation_args, cval=nan)
-
-    # If a moving mask is provided, transform this as well
-    if moving_mask is not None:
-        if transformation_args is None:
-            transformed_mask = transformation(moving_mask, moving_affine, static.shape, static_affine, params)
-        else:
-            transformed_mask = transformation(moving_mask, moving_affine, static.shape, static_affine, params, transformation_args)
-        transformed_mask[transformed_mask > 0.5] = 1
-        transformed_mask[transformed_mask <= 0.5] = 0
-
-    # Get the indices in the static data where the mask is non-zero
-    if static_mask is not None and moving_mask is not None:
-        ind = np.where(transformed!=nan and static_mask==1 and transformed_mask==1)
-    elif static_mask is not None:
-        ind = np.where(transformed!=nan and static_mask==1) 
-    elif moving_mask is not None:
-        ind = np.where(transformed!=nan and transformed_mask==1)
-    else:
-        ind = np.where(transformed!=nan)
-        
-    # Calculate the cost function         
-    if metric_args is None:
-        return metric(static[ind], transformed[ind])
+        st, tr = static[i], transformed[i]
     else:
-        return metric(static[ind], transformed[ind], metric_args)
+        st, tr = static, transformed
+    # Calculate 2d histogram
+    hist_2d, _, _ = np.histogram2d(st.ravel(), tr.ravel(), bins=20)
+    # Convert bins counts to probability values
+    pxy = hist_2d / float(np.sum(hist_2d))
+    px = np.sum(pxy, axis=1) # marginal for x over y
+    py = np.sum(pxy, axis=0) # marginal for y over x
+    px_py = px[:, None] * py[None, :] # Broadcast to multiply marginals
+    # Now we can do the calculation using the pxy, px_py 2D arrays
+    nzs = pxy > 0 # Only non-zero pxy values contribute to the sum
+    return -np.sum(pxy[nzs] * np.log(pxy[nzs] / px_py[nzs]))
 
 
 def print_current(xk):
     print('Current parameter estimates: ' , xk)
     return False
 
+
 def print_current_norm(xk):
     print('Norm of current parameter estimates: ' , np.linalg.norm(xk))
     return False
@@ -892,10 +991,28 @@ def minimize(*args, method='GD', **kwargs):
 
     if method == 'GD':
         return minimize_gd(*args, **kwargs)
+    elif method=='brute':
+        return minimize_brute(*args, **kwargs)
     else:
         res = opt.minimize(*args, method=method, **kwargs)
         return res.x
-    
+
+
+def minimize_brute(cost_function, parameters, args=None, callback=None, options={}):
+    #options = {'grid':[[start, stop, num], [start, stop, num], ...]}
+    x = [np.linspace(p[0], p[1], p[2]) for p in options['grid']]
+    x = np.meshgrid(*tuple(x), indexing='ij')
+    for i in range(x[0].size):
+        parameters = np.array([xp.ravel()[i] for xp in x])
+        cost = cost_function(parameters, *args)
+        if i==0:
+            minimum = cost
+            par_opt = parameters
+        elif cost<minimum:
+            minimum = cost
+            par_opt = parameters
+    return par_opt
+
 
 def minimize_gd(cost_function, parameters, args=None, callback=None, options={}, bounds=None):
 
@@ -931,7 +1048,6 @@ def minimize_gd(cost_function, parameters, args=None, callback=None, options={},
             return parameters
 
 
-
 def gradient(cost_function, parameters, f0, step, bounds, *args):
     grad = np.empty(parameters.shape)
     for i in range(parameters.size):
@@ -952,8 +1068,7 @@ def gradient(cost_function, parameters, f0, step, bounds, *args):
     # Normalize the gradient
     grad_norm = np.linalg.norm(grad)
     if grad_norm == 0:
-        msg = 'Gradient has length zero - cannot perform gradient descent'
-        raise ValueError(msg) 
+        return grad 
     grad /= grad_norm
     grad = np.multiply(step, grad)
 
@@ -1044,18 +1159,31 @@ def line_search(cost_function, grad, p0, stepsize0, f0, bounds, *args, tolerance
             raise ValueError(msg) 
         
 
-def goodness_of_alignment(params, transformation, metric, moving, moving_affine, static, static_affine, coord):
+def goodness_of_alignment(params, transformation, metric, moving, moving_affine, static, static_affine, coord, moving_mask, static_mask_ind):
 
     # Transform the moving image
     nan = 2**16-2 #np.nan does not work
     transformed = transformation(moving, moving_affine, static.shape, static_affine, params, output_coordinates=coord, cval=nan)
-        
-    # Calculate the cost function  
-    ls = metric(static, transformed, nan=nan)
     
-    return ls
+    # If a moving mask is provided, this needs to be transformed in the same way
+    if moving_mask is None:
+        moving_mask_ind = None
+    else:
+        mask_transformed = transformation(moving_mask, moving_affine, static.shape, static_affine, params, output_coordinates=coord, cval=nan)
+        moving_mask_ind = np.where(mask_transformed >= 0.5)
+
+    # Calculate matric in indices exposed by the mask(s)
+    if static_mask_ind is None and moving_mask_ind is None:
+        return metric(static, transformed, nan=nan)
+    if static_mask_ind is None and moving_mask_ind is not None:
+        return metric(static[moving_mask_ind], transformed[moving_mask_ind], nan=nan)
+    if static_mask_ind is not None and moving_mask_ind is None:
+        return metric(static[static_mask_ind], transformed[static_mask_ind], nan=nan)
+    if static_mask_ind is not None and moving_mask_ind is not None:
+        ind = static_mask_ind or moving_mask_ind
+        return metric(static[ind], transformed[ind], nan=nan)
+    
     
-
 def align(
         moving = None, 
         static = None, 
@@ -1065,7 +1193,11 @@ def align(
         transformation = translate,
         metric = sum_of_squares,
         optimization = {'method':'GD', 'options':{}},
-        resolutions = [1]):
+        resolutions = [1], 
+        static_mask = None,
+        static_mask_affine = None, 
+        moving_mask = None,
+        moving_mask_affine = None):
     
     # Set defaults
     if moving is None:
@@ -1074,9 +1206,6 @@ def align(
     if static is None:
         msg = 'The static volume is a required argument for alignment'
         raise ValueError(msg)
-    if parameters is None:
-        msg = 'Initial values for alignment must be provided'
-        raise ValueError(msg)
     if moving.ndim == 2: # If 2d array, add a 3d dimension of size 1
         moving = np.expand_dims(moving, axis=-1)
     if static.ndim == 2: # If 2d array, add a 3d dimension of size 1
@@ -1085,7 +1214,12 @@ def align(
         moving_affine = np.eye(1 + moving.ndim)
     if static_affine is None:
         static_affine = np.eye(1 + static.ndim)
-
+    if moving_mask is not None:
+        if moving_mask_affine is None:
+            moving_mask_affine = moving_affine
+    if static_mask is not None:
+        if static_mask_affine is None:
+            static_mask_affine = static_affine
 
     # Perform multi-resolution loop
     for res in resolutions:
@@ -1106,19 +1240,27 @@ def align(
             static_resampled_affine = affine_matrix(rotation=r, translation=t, pixel_spacing=p*res)
             static_resampled, static_resampled_affine = affine_reslice(static, static_affine, static_resampled_affine)
 
+        # resample the masks on the geometry of the target volumes
+        if moving_mask is None:
+            moving_mask_resampled = None
+        else:
+            moving_mask_resampled, _ = affine_reslice(moving_mask, moving_mask_affine, moving_resampled_affine, moving_resampled.shape)
+        if static_mask is None:
+            static_mask_resampled_ind = None
+        else:
+            static_mask_resampled, _ = affine_reslice(static_mask, static_mask_affine, static_resampled_affine, static_resampled.shape)
+            static_mask_resampled_ind = np.where(static_mask_resampled >= 0.5)
+
         coord = volume_coordinates(static_resampled.shape) 
         # Here we need a generic precomputation step:
         # prec = precompute(moving_resampled, moving_resampled_affine, static_resampled, static_resampled_affine)
         # args = (transformation, metric, moving_resampled, moving_resampled_affine, static_resampled, static_resampled_affine, coord, prec)
-        args = (transformation, metric, moving_resampled, moving_resampled_affine, static_resampled, static_resampled_affine, coord)
+        args = (transformation, metric, moving_resampled, moving_resampled_affine, static_resampled, static_resampled_affine, coord, moving_mask_resampled, static_mask_resampled_ind)
         parameters = minimize(goodness_of_alignment, parameters, args=args, **optimization)
 
     return parameters
 
 
-
-
-
 def align_slice_by_slice(
         moving = None, 
         static = None, 
@@ -1130,22 +1272,34 @@ def align_slice_by_slice(
         optimization = {'method':'GD', 'options':{}},
         resolutions = [1],
         slice_thickness = None,
-        progress = None):
+        progress = None,
+        static_mask = None,
+        static_mask_affine = None, 
+        moving_mask = None,
+        moving_mask_affine = None):
     
     # If a single slice thickness is provided, turn it into a list.
     nz = moving.shape[2]
     if slice_thickness is not None:
         if not isinstance(slice_thickness, list):
             slice_thickness = [slice_thickness]*nz
+    if not isinstance(parameters, list):
+        parameters = [parameters]*nz
 
     estimate = []
     for z in range(nz):
 
+        print('SLICE: ', z)
+
         if progress is not None:
             progress(z, nz)
         
         # Get the slice and its affine
         moving_z, moving_affine_z = extract_slice(moving, moving_affine, z, slice_thickness)
+        if moving_mask is None:
+            moving_mask_z, moving_mask_affine_z = None, None
+        else:
+            moving_mask_z, moving_mask_affine_z = extract_slice(moving_mask, moving_mask_affine, z, slice_thickness)
 
         # Align volumes
         try:
@@ -1154,14 +1308,18 @@ def align_slice_by_slice(
                 moving_affine = moving_affine_z, 
                 static = static, 
                 static_affine = static_affine, 
-                parameters = parameters, 
+                parameters = parameters[z], 
                 resolutions = resolutions, 
                 transformation = transformation,
                 metric = metric,
                 optimization = optimization,
+                static_mask = static_mask,
+                static_mask_affine = static_mask_affine, 
+                moving_mask = moving_mask_z,
+                moving_mask_affine = moving_mask_affine_z,
             )
         except:
-            estimate_z = parameters
+            estimate_z = parameters[z]
 
         estimate.append(estimate_z)
 
@@ -1169,8 +1327,6 @@ def align_slice_by_slice(
 
 
 
-
-
 #############################
 # Plot results
 #############################
@@ -1272,6 +1428,21 @@ def plot_affine_transformed(input_data, input_affine, output_data, output_affine
     pl.show()
 
 
+def plot_bounding_box(input_data, input_affine, output_shape, output_affine):
+
+    pl = plot_volume(input_data, input_affine)
+
+    # Plot wireframe around edges of resliced volume
+    vertices, faces = parallellepid(output_shape, affine=output_affine)
+    box = pv.PolyData(vertices, faces)
+    pl.add_mesh(box,
+        style='wireframe',
+        opacity=1.0,
+        color=(255,0,0), 
+    )
+    pl.show()
+
+
 def plot_freeform_transformed(input_data, input_affine, output_data, output_affine):
 
     pl = plot_affine_resliced(input_data, input_affine, output_data, output_affine)
@@ -1711,6 +1882,29 @@ def test_plot(n=1):
     pl.show()
 
 
+def test_bounding_box():
+
+    # Define geometry of source data
+    input_shape = np.array([300, 250, 12])   # mm
+    pixel_spacing = np.array([1.25, 1.25, 10]) # mm
+    rotation_angle = 0.5 * (np.pi/2) # radians
+    rotation_axis = [1,0,0]
+    translation = np.array([0, -40, 180]) # mm
+    margin = 10 # mm
+
+    # Generate reference volume
+    rotation = rotation_angle * np.array(rotation_axis)/np.linalg.norm(rotation_axis)
+    input_affine = affine_matrix(rotation=rotation, translation=translation, pixel_spacing=pixel_spacing)
+    #input_data, input_affine = generate('double ellipsoid', shape=input_shape, affine=input_affine, markers=False)
+    input_data, input_affine = generate('triple ellipsoid', shape=input_shape, affine=input_affine, markers=False)
+
+    # Perform translation with reshaping
+    output_data, output_affine = bounding_box(input_data, input_affine, margin=margin)
+
+    # Display results
+    plot_bounding_box(input_data, input_affine, output_data.shape, output_affine)
+
+
 def test_affine_reslice(n=1):
 
     if n==1:
@@ -2528,7 +2722,6 @@ def test_align_freeform(n=1):
     pl.show()
 
 
-
 def test_align_freeform_pyramid(n=1):
 
     if n==1:
@@ -2580,14 +2773,13 @@ def test_align_freeform_pyramid(n=1):
         
 
 
-
 if __name__ == "__main__":
-
-    dataset=2
+    pass
 
     # Test plotting
     # -------------
     # test_plot(1)
+    # test_bounding_box()
 
 
     # Test affine transformations
@@ -2621,6 +2813,6 @@ if __name__ == "__main__":
     # test_align_stretch(n=2)
     # test_align_rigid(n=1)
     # test_align_affine(n=1)
-    test_align_freeform(n=3)
+    # test_align_freeform(n=3)