From b7dd98bdbd26dc1d00efb6fd5fc818c9bb4fe7bf Mon Sep 17 00:00:00 2001 From: Oscar Esteban Date: Wed, 16 Nov 2022 12:10:17 +0100 Subject: [PATCH 1/6] enh(BSplineApprox): integrate downsampling when the input is high-res Adds a new input ``zooms_min`` that will trigger the downsampling of the input field map when it comes with a very high resolution (typically this is the case for SDC SyN). Although the input is resampled, the output field is interpolated at the original resolution for debugging purposes. Downsampling dramatically speeds up approximation. Alternatively, a solution (for the SDC SyN) would be to calculate a heavily dilated mask of the brain, and use it to limit the number of voxels that are fit in regression. --- sdcflows/interfaces/bspline.py | 130 +++++++++++++++++++++++---------- 1 file changed, 93 insertions(+), 37 deletions(-) diff --git a/sdcflows/interfaces/bspline.py b/sdcflows/interfaces/bspline.py index 46c27db90e..74af8be549 100644 --- a/sdcflows/interfaces/bspline.py +++ b/sdcflows/interfaces/bspline.py @@ -47,7 +47,6 @@ DEFAULT_ZOOMS_MM = (40.0, 40.0, 20.0) # For human adults (mid-frequency), in mm DEFAULT_LF_ZOOMS_MM = (100.0, 100.0, 40.0) # For human adults (low-frequency), in mm DEFAULT_HF_ZOOMS_MM = (16.0, 16.0, 10.0) # For human adults (high-frequency), in mm -BSPLINE_SUPPORT = 2 - 1.82e-3 # Disallows weights < 1e-9 LOGGER = logging.getLogger("nipype.interface") @@ -76,6 +75,11 @@ class _BSplineApproxInputSpec(BaseInterfaceInputSpec): usedefault=True, desc="generate a field, extrapolated outside the brain mask", ) + zooms_min = traits.Float( + 1.4, + usedefault=True, + desc="limit minimum image zooms, set 0.0 to use the original image", + ) class _BSplineApproxOutputSpec(TraitedSpec): @@ -123,16 +127,45 @@ class BSplineApprox(SimpleInterface): def _run_interface(self, runtime): from sklearn import linear_model as lm - from scipy.sparse import vstack as sparse_vstack + + # Output name baseline + out_name = fname_presuffix( + self.inputs.in_data, suffix="_field", newpath=runtime.cwd + ) # Load in the fieldmap fmapnii = nb.load(self.inputs.in_data) + zooms = fmapnii.header.get_zooms() + + # Get a mask (or define on the spot to cover the full extent) + masknii = ( + nb.load(self.inputs.in_mask) + if isdefined(self.inputs.in_mask) + else None + ) + + need_resize = np.any(np.array(zooms) < self.inputs.zooms_min) + if need_resize: + from niworkflows.utils.images import resample_by_spacing + + LOGGER.info( + "Resampling image with resolution exceeding 'zooms_min' " + f"({'x'.join(str(s) for s in zooms)})." + ) + fmapnii = resample_by_spacing(fmapnii, [self.inputs.zooms_min] * 3) + + if masknii is not None: + masknii = resample_by_spacing(masknii, [self.inputs.zooms_min] * 3) + data = fmapnii.get_fdata(dtype="float32") + + # Generate a numpy array with the mask mask = ( - nb.load(self.inputs.in_mask).get_fdata() > 0 - if isdefined(self.inputs.in_mask) - else np.ones_like(data, dtype=bool) + np.ones_like(fmapnii.dataobj, dtype=bool) if masknii is None + else np.asanyarray(masknii.dataobj) > 1e-4 ) + + # Convert spacings to numpy arrays bs_spacing = [np.array(sp, dtype="float32") for sp in self.inputs.bs_spacing] # Recenter the fieldmap @@ -145,45 +178,29 @@ def _run_interface(self, runtime): elif self.inputs.recenter == "mean": data -= np.mean(data[mask]) - # Calculate the spatial location of control points - bs_levels = [] - ncoeff = [] - weights = None - for sp in bs_spacing: - level = bspline_grid(fmapnii, control_zooms_mm=sp) - bs_levels.append(level) - ncoeff.append(level.dataobj.size) - - weights = ( - gbsw(fmapnii, level) - if weights is None - else sparse_vstack((weights, gbsw(fmapnii, level))) - ) + # Calculate collocation matrix & the spatial location of control points + colmat, bs_levels = _collocation_matrix(fmapnii, bs_spacing) - regressors = weights.T.tocsr()[mask.reshape(-1), :] + bs_levels_str = ['x'.join(str(s) for s in level.shape) for level in bs_levels] + bs_levels_str[-1] = f"and {bs_levels_str[-1]}" + LOGGER.info( + f"Approximating B-Splines grids ({', '.join(bs_levels_str)} [knots]) on a grid of " + f"{'x'.join(str(s) for s in fmapnii.shape)} ({np.prod(fmapnii.shape)}) voxels," + f" of which {mask.sum()} fall within the mask." + ) # Fit the model model = lm.Ridge(alpha=self.inputs.ridge_alpha, fit_intercept=False) - model.fit(regressors, data[mask]) - - interp_data = np.zeros_like(data) - interp_data[mask] = np.array(model.coef_) @ regressors.T # Interpolation - - # Store outputs - out_name = fname_presuffix( - self.inputs.in_data, suffix="_field", newpath=runtime.cwd - ) - hdr = fmapnii.header.copy() - hdr.set_data_dtype("float32") - fmapnii.__class__(interp_data, fmapnii.affine, hdr).to_filename(out_name) - self._results["out_field"] = out_name + model.fit(colmat[mask.reshape(-1), :], data[mask]) + # Store coefficients index = 0 self._results["out_coeff"] = [] - for i, (n, bsl) in enumerate(zip(ncoeff, bs_levels)): + for i, bsl in enumerate(bs_levels): + n = bsl.dataobj.size out_level = out_name.replace("_field.", f"_coeff{i:03}.") bsl.__class__( - np.array(model.coef_, dtype="float32")[index : index + n].reshape( + np.array(model.coef_, dtype="float32")[index:index + n].reshape( bsl.shape ), bsl.affine, @@ -192,6 +209,27 @@ def _run_interface(self, runtime): index += n self._results["out_coeff"].append(out_level) + # Interpolating in the original grid will require a new collocation matrix + if need_resize: + fmapnii = nb.load(self.inputs.in_data) + data = fmapnii.get_fdata(dtype="float32") + mask = ( + np.ones_like(fmapnii.dataobj, dtype=bool) if masknii is None + else np.asanyarray(nb.load(self.inputs.in_mask).dataobj) > 1e-4 + ) + colmat, _ = _collocation_matrix(fmapnii, bs_spacing) + + regressors = colmat[mask.reshape(-1), :] + interp_data = np.zeros_like(data) + # Interpolate the field from the coefficients just calculated + interp_data[mask] = regressors @ model.coef_ + + # Store interpolated field + hdr = fmapnii.header.copy() + hdr.set_data_dtype("float32") + fmapnii.__class__(interp_data, fmapnii.affine, hdr).to_filename(out_name) + self._results["out_field"] = out_name + # Write out fitting-error map self._results["out_error"] = out_name.replace("_field.", "_error.") fmapnii.__class__( @@ -205,8 +243,8 @@ def _run_interface(self, runtime): self._results["out_extrapolated"] = self._results["out_field"] return runtime - extrapolators = weights.tocsc()[:, ~mask.reshape(-1)] - interp_data[~mask] = np.array(model.coef_) @ extrapolators # Extrapolation + extrapolators = colmat[~mask.reshape(-1), :] + interp_data[~mask] = extrapolators @ model.coef_ # Extrapolation self._results["out_extrapolated"] = out_name.replace("_field.", "_extra.") fmapnii.__class__(interp_data, fmapnii.affine, hdr).to_filename( self._results["out_extrapolated"] @@ -457,6 +495,24 @@ def bspline_grid(img, control_zooms_mm=DEFAULT_ZOOMS_MM): return img.__class__(np.zeros(bs_shape, dtype="float32"), bs_affine) +def _collocation_matrix(image, knot_spacing): + from scipy.sparse import vstack as sparse_vstack + + bs_levels = [] + weights = None + for sp in knot_spacing: + level = bspline_grid(image, control_zooms_mm=sp) + bs_levels.append(level) + + weights = ( + gbsw(image, level) + if weights is None + else sparse_vstack((weights, gbsw(image, level))) + ) + + return weights.T.tocsr(), bs_levels + + def _fix_topup_fieldcoeff(in_coeff, fmap_ref, pe_dir, out_file=None): """Read in a coefficients file generated by TOPUP and fix x-form headers.""" from pathlib import Path From 9893475799654274fbfc8c35b060a1453d51846c Mon Sep 17 00:00:00 2001 From: Oscar Esteban Date: Thu, 17 Nov 2022 09:23:08 +0100 Subject: [PATCH 2/6] enh: allow tuples for the ``zooms_min`` input Allows anisotropic voxel sizes when resampling the input data. cc/ @eilidhmacnicol Co-authored-by: Chris Markiewicz --- sdcflows/interfaces/bspline.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/sdcflows/interfaces/bspline.py b/sdcflows/interfaces/bspline.py index 74af8be549..0adcf13ae2 100644 --- a/sdcflows/interfaces/bspline.py +++ b/sdcflows/interfaces/bspline.py @@ -23,6 +23,7 @@ """Filtering of :math:`B_0` field mappings with B-Splines.""" from itertools import product from pathlib import Path +from contextlib import suppress import numpy as np import nibabel as nb from nibabel.affines import apply_affine @@ -75,8 +76,10 @@ class _BSplineApproxInputSpec(BaseInterfaceInputSpec): usedefault=True, desc="generate a field, extrapolated outside the brain mask", ) - zooms_min = traits.Float( - 1.4, + zooms_min = traits.Union( + traits.Float, + traits.Tuple(traits.Float, traits.Float, traits.Float), + default_value=1.95, usedefault=True, desc="limit minimum image zooms, set 0.0 to use the original image", ) @@ -148,14 +151,19 @@ def _run_interface(self, runtime): if need_resize: from niworkflows.utils.images import resample_by_spacing + zooms_min = self.inputs.zooms_min + + with suppress(TypeError): + zooms_min = [float(zooms_min)] * 3 + LOGGER.info( "Resampling image with resolution exceeding 'zooms_min' " f"({'x'.join(str(s) for s in zooms)})." ) - fmapnii = resample_by_spacing(fmapnii, [self.inputs.zooms_min] * 3) + fmapnii = resample_by_spacing(fmapnii, zooms_min) if masknii is not None: - masknii = resample_by_spacing(masknii, [self.inputs.zooms_min] * 3) + masknii = resample_by_spacing(masknii, zooms_min) data = fmapnii.get_fdata(dtype="float32") From 4b0a53ccb2e84bee09ea46352eb84a7728185232 Mon Sep 17 00:00:00 2001 From: Oscar Esteban Date: Thu, 17 Nov 2022 10:08:03 +0100 Subject: [PATCH 3/6] fix: do not use ``niworkflows.utils.image.resample_by_spacing`` Replace the controverted function with a new implementation, which could be condiered as a candidate to replace Niworkflows'. Co-authored-by: Chris Markiewicz Co-authored-by: Mathias Goncalves --- sdcflows/interfaces/bspline.py | 6 ++--- sdcflows/utils/tools.py | 48 ++++++++++++++++++++++++++++++++++ 2 files changed, 51 insertions(+), 3 deletions(-) diff --git a/sdcflows/interfaces/bspline.py b/sdcflows/interfaces/bspline.py index 0adcf13ae2..424ece8bee 100644 --- a/sdcflows/interfaces/bspline.py +++ b/sdcflows/interfaces/bspline.py @@ -149,7 +149,7 @@ def _run_interface(self, runtime): need_resize = np.any(np.array(zooms) < self.inputs.zooms_min) if need_resize: - from niworkflows.utils.images import resample_by_spacing + from sdcflows.utils.tools import resample_to_zooms zooms_min = self.inputs.zooms_min @@ -160,10 +160,10 @@ def _run_interface(self, runtime): "Resampling image with resolution exceeding 'zooms_min' " f"({'x'.join(str(s) for s in zooms)})." ) - fmapnii = resample_by_spacing(fmapnii, zooms_min) + fmapnii = resample_to_zooms(fmapnii, zooms_min) if masknii is not None: - masknii = resample_by_spacing(masknii, zooms_min) + masknii = resample_to_zooms(masknii, zooms_min) data = fmapnii.get_fdata(dtype="float32") diff --git a/sdcflows/utils/tools.py b/sdcflows/utils/tools.py index 4942f6a05e..1edcb02a33 100644 --- a/sdcflows/utils/tools.py +++ b/sdcflows/utils/tools.py @@ -23,6 +23,54 @@ """Image processing tools.""" +def resample_to_zooms(in_file, zooms, order=3, prefilter=True): + """Resample the input data to a new grid with the requested zooms.""" + from pathlib import Path + import numpy as np + import nibabel as nb + from nibabel.affines import apply_affine + from nitransforms.linear import Affine + + if isinstance(in_file, (str, Path)): + in_file = nb.load(in_file) + + # Prepare output x-forms + sform, scode = in_file.get_sform(coded=True) + qform, qcode = in_file.get_qform(coded=True) + + hdr = in_file.header.copy() + zooms = np.array(zooms) + + # Calculate the factors to normalize voxel size to the specific zooms + pre_zooms = np.array(in_file.header.get_zooms()[:3]) + pre_shape = np.array(in_file.shape[:3]) + pre_affine = in_file.affine + + # Calculate new affine and shape + affine = np.eye(4) + affine[:-1, :-1] = zooms * pre_affine[:-1, :-1] / pre_zooms + + extremes = apply_affine(pre_affine, [(0, 0, 0), pre_shape - 1]) + extent = np.abs(extremes[1] - extremes[0]) + + new_shape = np.ceil(extent / zooms).astype(int) + affine[:-1, -1] = apply_affine(pre_affine, (pre_shape - 1) * 0.5) - affine[ + :-1, :-1 + ] @ ((new_shape - 1) * 0.5) + + # Generate new reference + hdr.set_sform(affine, scode) + hdr.set_qform(affine, qcode) + newref = in_file.__class__( + np.zeros(new_shape, dtype=hdr.get_data_dtype()), + affine, + hdr, + ) + + # Resample via identity transform + return Affine(reference=newref).apply(in_file, order=order, prefilter=prefilter) + + def ensure_positive_cosines(img): """ Reorient axes polarity to have all positive direction cosines. From 41c57887661249efa15bb11a7b4698c17fab2862 Mon Sep 17 00:00:00 2001 From: Oscar Esteban Date: Thu, 17 Nov 2022 10:14:16 +0100 Subject: [PATCH 4/6] enh: remove function alias Co-authored-by: Mathias Goncalves --- sdcflows/interfaces/bspline.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/sdcflows/interfaces/bspline.py b/sdcflows/interfaces/bspline.py index 424ece8bee..1972b8f446 100644 --- a/sdcflows/interfaces/bspline.py +++ b/sdcflows/interfaces/bspline.py @@ -41,7 +41,7 @@ OutputMultiObject, ) -from sdcflows.transform import grid_bspline_weights as gbsw +from sdcflows.transform import grid_bspline_weights LOW_MEM_BLOCK_SIZE = 1000 @@ -513,9 +513,9 @@ def _collocation_matrix(image, knot_spacing): bs_levels.append(level) weights = ( - gbsw(image, level) + grid_bspline_weights(image, level) if weights is None - else sparse_vstack((weights, gbsw(image, level))) + else sparse_vstack((weights, grid_bspline_weights(image, level))) ) return weights.T.tocsr(), bs_levels From 995a898f1097666a9438674d146097a87d905899 Mon Sep 17 00:00:00 2001 From: Oscar Esteban Date: Thu, 17 Nov 2022 17:29:15 +0100 Subject: [PATCH 5/6] Apply suggestions from code review Co-authored-by: Chris Markiewicz --- sdcflows/interfaces/bspline.py | 2 +- sdcflows/utils/tools.py | 20 ++++---------------- 2 files changed, 5 insertions(+), 17 deletions(-) diff --git a/sdcflows/interfaces/bspline.py b/sdcflows/interfaces/bspline.py index 846bbe3af6..7cdc1334eb 100644 --- a/sdcflows/interfaces/bspline.py +++ b/sdcflows/interfaces/bspline.py @@ -170,7 +170,7 @@ def _run_interface(self, runtime): # Generate a numpy array with the mask mask = ( - np.ones_like(fmapnii.dataobj, dtype=bool) if masknii is None + np.ones(fmapnii.shape, dtype=bool) if masknii is None else np.asanyarray(masknii.dataobj) > 1e-4 ) diff --git a/sdcflows/utils/tools.py b/sdcflows/utils/tools.py index 1edcb02a33..0bd816ebda 100644 --- a/sdcflows/utils/tools.py +++ b/sdcflows/utils/tools.py @@ -28,7 +28,7 @@ def resample_to_zooms(in_file, zooms, order=3, prefilter=True): from pathlib import Path import numpy as np import nibabel as nb - from nibabel.affines import apply_affine + from nibabel.affines import rescale_affine from nitransforms.linear import Affine if isinstance(in_file, (str, Path)): @@ -41,22 +41,10 @@ def resample_to_zooms(in_file, zooms, order=3, prefilter=True): hdr = in_file.header.copy() zooms = np.array(zooms) - # Calculate the factors to normalize voxel size to the specific zooms pre_zooms = np.array(in_file.header.get_zooms()[:3]) - pre_shape = np.array(in_file.shape[:3]) - pre_affine = in_file.affine - - # Calculate new affine and shape - affine = np.eye(4) - affine[:-1, :-1] = zooms * pre_affine[:-1, :-1] / pre_zooms - - extremes = apply_affine(pre_affine, [(0, 0, 0), pre_shape - 1]) - extent = np.abs(extremes[1] - extremes[0]) - - new_shape = np.ceil(extent / zooms).astype(int) - affine[:-1, -1] = apply_affine(pre_affine, (pre_shape - 1) * 0.5) - affine[ - :-1, :-1 - ] @ ((new_shape - 1) * 0.5) + # Could use `np.ceil` if we prefer + new_shape = np.rint(np.array(in_file.shape[:3]) * pre_zooms / zooms) + affine = rescale_affine(in_file.affine, in_file.shape[:3], zooms, new_shape) # Generate new reference hdr.set_sform(affine, scode) From 80bb7462f8472aaa478d471527a3371d3afeee62 Mon Sep 17 00:00:00 2001 From: Oscar Esteban Date: Thu, 17 Nov 2022 21:17:27 +0100 Subject: [PATCH 6/6] enh: final fixes Co-authored-by: Chris Markiewicz --- sdcflows/interfaces/bspline.py | 9 +++------ sdcflows/utils/tools.py | 2 +- 2 files changed, 4 insertions(+), 7 deletions(-) diff --git a/sdcflows/interfaces/bspline.py b/sdcflows/interfaces/bspline.py index 7cdc1334eb..31e65bd157 100644 --- a/sdcflows/interfaces/bspline.py +++ b/sdcflows/interfaces/bspline.py @@ -23,7 +23,6 @@ """Filtering of :math:`B_0` field mappings with B-Splines.""" from itertools import product from pathlib import Path -from contextlib import suppress import numpy as np import nibabel as nb from nibabel.affines import apply_affine @@ -152,14 +151,12 @@ def _run_interface(self, runtime): if need_resize: from sdcflows.utils.tools import resample_to_zooms - zooms_min = self.inputs.zooms_min - - with suppress(TypeError): - zooms_min = [float(zooms_min)] * 3 + zooms_min = np.maximum(zooms, self.inputs.zooms_min) LOGGER.info( "Resampling image with resolution exceeding 'zooms_min' " - f"({'x'.join(str(s) for s in zooms)})." + f"({'x'.join(str(s) for s in zooms)} → " + f"{'x'.join(str(s) for s in zooms_min)})." ) fmapnii = resample_to_zooms(fmapnii, zooms_min) diff --git a/sdcflows/utils/tools.py b/sdcflows/utils/tools.py index 0bd816ebda..5fe23051eb 100644 --- a/sdcflows/utils/tools.py +++ b/sdcflows/utils/tools.py @@ -50,7 +50,7 @@ def resample_to_zooms(in_file, zooms, order=3, prefilter=True): hdr.set_sform(affine, scode) hdr.set_qform(affine, qcode) newref = in_file.__class__( - np.zeros(new_shape, dtype=hdr.get_data_dtype()), + np.zeros(new_shape.astype(int), dtype=hdr.get_data_dtype()), affine, hdr, )