ENH: Add masked/bspline fitting variant of Nyul histogram matching.

ants/utils/__init__.py

@@ -16,7 +16,7 @@
 from .get_mask import get_mask
 from .get_neighborhood import (get_neighborhood_in_mask,
  get_neighborhood_at_voxel)
-from .histogram_match_image import histogram_match_image
+from .histogram_match_image import histogram_match_image, histogram_match_image2
 from .histogram_equalize_image import histogram_equalize_image
 from .hausdorff_distance import hausdorff_distance
 from .image_similarity import image_similarity

ants/utils/histogram_match_image.py

@@ -1,11 +1,14 @@
 
-__all__ = ['histogram_match_image']
+__all__ = ['histogram_match_image',
+ 'histogram_match_image2']
 
-import math
+import numpy as np
 
-from ..core import ants_image as iio
+from ..core import ants_image_io as iio
 from .. import utils
 
+from ..utils import fit_bspline_object_to_scattered_data
+
 
 def histogram_match_image(source_image, reference_image, number_of_histogram_bins=255, number_of_match_points=64, use_threshold_at_mean_intensity=False):
  """
@@ -51,3 +54,102 @@ def histogram_match_image(source_image, reference_image, number_of_histogram_bin
  return new_image
 
 
+def histogram_match_image2(source_image, reference_image, 
+ source_mask=None, reference_mask=None,
+ match_points=64,
+ transform_domain_size=255):
+ """
+ Transform image intensities based on histogram mapping.
+
+ Apply B-spline 1-D maps to an input image for intensity warping.
+
+ Arguments
+ ---------
+ source_image : ANTsImage
+ source image
+
+ reference_image : ANTsImage
+ reference image
+
+ source_mask : ANTsImage
+ source mask
+
+ reference_mask : ANTsImage
+ reference mask
+
+ match_points : integer or tuple
+ Parametric points at which the intensity transform displacements are 
+ specified between [0, 1], i.e. quantiles. Alternatively, a single number 
+ can be given and the sequence is linearly spaced in [0, 1]. 
+
+ transform_domain_size : integer
+ Defines the sampling resolution of the B-spline warping.
+
+ Returns
+ -------
+ ANTs image
+
+ Example
+ -------
+ >>> import ants
+ >>> src_img = ants.image_read(ants.get_data('r16'))
+ >>> ref_img = ants.image_read(ants.get_data('r64'))
+ >>> src_ref = ants.histogram_match_image(src_img, ref_img)
+ """
+
+ if not isinstance(match_points, int):
+ if any(b < 0 for b in match_points) and any(b > 1 for b in match_points):
+ raise ValueError("If specifying match_points as a vector, values must be in the range [0, 1]")
+
+ # Use entire image if mask isn't specified
+ if source_mask is None:
+ source_mask = source_image * 0 + 1
+ if reference_mask is None:
+ reference_mask = reference_image * 0 + 1
+
+ source_array = source_image.numpy()
+ source_mask_array = source_mask.numpy()
+ source_masked_min = source_image[source_mask != 0].min()
+ source_masked_max = source_image[source_mask != 0].max()
+
+ reference_array = reference_image.numpy()
+ reference_mask_array = reference_mask.numpy()
+
+ parametric_points = None
+ if not isinstance(match_points, int):
+ parametric_points = match_points
+ else:
+ parametric_points = np.linspace(0, 1, match_points)
+
+ source_intensity_quantiles = np.quantile(source_array[source_mask_array != 0], parametric_points)
+ reference_intensity_quantiles = np.quantile(reference_array[reference_mask_array != 0], parametric_points)
+ displacements = reference_intensity_quantiles - source_intensity_quantiles
+
+ scattered_data = np.reshape(displacements, (len(displacements), 1))
+ parametric_data = np.reshape(parametric_points * (source_masked_max - source_masked_min) + source_masked_min, (len(parametric_points), 1))
+
+ transform_domain_origin = source_masked_min
+ transform_domain_spacing = (source_masked_max - transform_domain_origin) / (transform_domain_size - 1)
+
+ bspline_histogram_transform = fit_bspline_object_to_scattered_data(scattered_data,
+ parametric_data, [transform_domain_origin], [transform_domain_spacing], [transform_domain_size],
+ data_weights=None, is_parametric_dimension_closed=None, number_of_fitting_levels=8, 
+ mesh_size=1, spline_order=3)
+
+ transform_domain = np.linspace(source_masked_min, source_masked_max, transform_domain_size)
+
+ transformed_source_array = source_image.numpy()
+ for i in range(len(transform_domain) - 1):
+ indices = np.where((source_array >= transform_domain[i]) & (source_array < transform_domain[i+1]))
+ intensities = source_array[indices]
+
+ alpha = (intensities - transform_domain[i])/(transform_domain[i+1] - transform_domain[i])
+ xfrm = alpha * (bspline_histogram_transform[i+1] - bspline_histogram_transform[i]) + bspline_histogram_transform[i]
+ transformed_source_array[indices] = intensities + xfrm
+
+ transformed_source_image = iio.from_numpy(transformed_source_array, origin=source_image.origin,
+ spacing=source_image.spacing, direction=source_image.direction)
+ transformed_source_image[source_mask == 0] = source_image[source_mask == 0]
+
+ return(transformed_source_image)
+