enh: add downsample() in addition to decimate()

Author: oesteban

src/nifreeze/data/filtering.py

@@ -28,6 +28,7 @@
 
 import numpy as np
 from nibabel import Nifti1Image, load
+from nibabel.affines import apply_affine, voxel_sizes
 from scipy.ndimage import gaussian_filter as _gs
 from scipy.ndimage import map_coordinates, median_filter
 from skimage.morphology import ball
@@ -141,15 +142,15 @@ def gaussian_filter(
     return _gs(data, vox_width)
 
 
-def decimate(
+def downsample(
     in_file: str,
-    factor: int | tuple[int, int, int],
+    shape: tuple[int, int, int],
     smooth: bool | tuple[int, int, int] = True,
     order: int = 3,
     nonnegative: bool = True,
 ) -> Nifti1Image:
     """
-    Decimates a 3D or 4D Nifti image by a specified downsampling factor.
+    Downsamples a 3D or 4D Nifti image by a specified downsampling factor.
 
     This function downsamples a Nifti image by averaging voxels within a user-defined
     factor in each spatial dimension. It optionally applies Gaussian smoothing
@@ -186,47 +187,54 @@ def decimate(
 
     imnii = load(in_file)
     data = np.squeeze(imnii.get_fdata())  # Remove unused dimensions
-    datashape = data.shape
+    datashape = np.array(data.shape)
+    shape = np.array(shape)
     ndim = data.ndim
 
-    if isinstance(factor, Number):
-        factor = tuple([factor] * min(3, ndim))
-
-    if any(f <= 0 for f in factor[:3]):
-        raise ValueError("All spatial downsampling factors must be positive.")
-
-    if ndim == 4 and len(factor) == 3:
-        factor = (*factor, 0)
-
     if smooth:
         if smooth is True:
-            smooth = factor
+            smooth = datashape[:3] / shape[:3]
         data = gaussian_filter(data, smooth)
 
+    extents = (
+        apply_affine(imnii.affine, datashape - 0.5)
+        - apply_affine(imnii.affine, (-0.5, -0.5, -0.5))
+    )
+    newzooms = extents / shape
+
+    # Update affine transformation
+    newaffine = np.eye(4)
+    oldzooms = voxel_sizes(imnii.affine)
+    newaffine[:3, :3] = np.diag(newzooms / oldzooms) @ imnii.affine[:3, :3]
+
+    # Update offset so new array is aligned with original
+    newaffine[:3, 3] = (
+        apply_affine(imnii.affine, 0.5 * datashape)
+        - apply_affine(newaffine, 0.5 * shape)
+    )
+
+    xfm = np.linalg.inv(imnii.affine) @ newaffine
+
     # Create downsampled grid
     down_grid = np.array(
         np.meshgrid(
-            *[np.arange(_s, step=int(_f) or 1) for _s, _f in zip(datashape, factor)],
+            *[np.arange(_s, step=1) for _s in shape],
             indexing="ij",
         )
     )
-    new_shape = down_grid.shape[1:]
-
-    # Update affine transformation
-    newaffine = imnii.affine.copy()
-    newaffine[:3, :3] = np.array(factor[:3]) * newaffine[:3, :3]
 
-    # TODO: Update offset so new array is aligned with original
+    # Locations is an Nx3 array of index coordinates of the original image where we sample
+    locations = apply_affine(xfm, down_grid.reshape((ndim, np.prod(shape))).T)
 
     # Resample data on the new grid
     resampled = map_coordinates(
         data,
-        down_grid.reshape((ndim, np.prod(new_shape))),
+        locations.T,
         order=order,
         mode="constant",
         cval=0,
         prefilter=True,
-    ).reshape(new_shape)
+    ).reshape(shape)
 
     # Set negative values to zero (optional)
     if order > 2 and nonnegative:
@@ -238,3 +246,80 @@ def decimate(
     newnii.set_qform(newaffine, code=1)
 
     return newnii
+
+
+def decimate(
+    in_file: str,
+    factor: int | tuple[int, int, int],
+    smooth: bool | tuple[int, int, int] = True,
+    order: int = 3,
+    nonnegative: bool = True,
+) -> Nifti1Image:
+    """
+    Decimates a 3D or 4D Nifti image by a specified downsampling factor.
+
+    This function downsamples a Nifti image by averaging voxels within a user-defined
+    factor in each spatial dimension. It optionally applies Gaussian smoothing
+    before downsampling to reduce aliasing artifacts. The function also handles
+    updating the affine transformation matrix to reflect the change in voxel size.
+
+    Parameters
+    ----------
+    in_file : :obj:`str`
+        Path to the input NIfTI image file.
+    factor : :obj:`int` or :obj:`tuple`
+        The downsampling factor. If a single integer is provided, it is applied
+        uniformly across all spatial dimensions. Alternatively, a tuple of three
+        integers can be provided to specify different downsampling factors for each
+        spatial dimension (x, y, z). Values must be greater than 0.
+    smooth : :obj:`bool` or :obj:`tuple`, optional (default=``True``)
+        Controls application of Gaussian smoothing before downsampling. If True,
+        a smoothing kernel size equal to the downsampling factor is applied.
+        Alternatively, a tuple of three integers can be provided to specify
+        different smoothing kernel sizes for each spatial dimension. Setting to
+        False disables smoothing.
+    order : :obj:`int`, optional (default=3)
+        The order of the spline interpolation used for downsampling. Higher
+        orders provide smoother results but are computationally more expensive.
+    nonnegative : :obj:`bool`, optional (default=``True``)
+        If True, negative values in the downsampled data are set to zero.
+
+    Returns
+    -------
+    :obj:`~nibabel.Nifti1Image`
+        The downsampled NIfTI image object.
+
+    """
+
+    imnii = load(in_file)
+    data = np.squeeze(imnii.get_fdata())  # Remove unused dimensions
+    ndim = data.ndim
+
+    if isinstance(factor, Number):
+        factor = tuple([factor] * min(3, ndim))
+
+    if any(f <= 0 for f in factor[:3]):
+        raise ValueError("All spatial downsampling factors must be positive.")
+
+    if ndim == 4 and len(factor) == 3:
+        factor = (*factor, 0)
+
+    if smooth:
+        if smooth is True:
+            smooth = factor
+        data = gaussian_filter(data, smooth)
+
+    # Update affine transformation
+    newaffine = imnii.affine.copy()
+    newaffine[:3, :3] = np.array(factor[:3]) * newaffine[:3, :3]
+
+    # Create new Nifti image with updated information
+    newnii = Nifti1Image(
+        data[::factor[0], ::factor[1], ::factor[2]],
+        newaffine,
+        imnii.header,
+    )
+    newnii.set_sform(newaffine, code=1)
+    newnii.set_qform(newaffine, code=1)
+
+    return newnii

test/test_filtering.py

@@ -26,23 +26,70 @@
 
 import pytest
 
-from nifreeze.data.filtering import decimate
+from nifreeze.data.filtering import decimate, downsample
 
 
 @pytest.mark.parametrize(
     ("size", "block_size"),
     [
-        ((20, 20, 20), (5, 5, 5),)
+        ((20, 20, 20), (5, 5, 5),),
+        ((21, 21, 21), (5, 5, 5),),
     ],
 )
-def test_decimation(tmp_path, size, block_size):
+@pytest.mark.parametrize(
+    ("zoom_x", ),
+    # [(1.0, ), (-1.0, ), (2.0, ), (-2.0, )],
+    [(2.0,)],
+)
+@pytest.mark.parametrize(
+    ("zoom_y", ),
+    # [(1.0, ), (-1.0, ), (2.0, ), (-2.0, )],
+    [(-2.0,)],
+)
+@pytest.mark.parametrize(
+    ("zoom_z", ),
+    # [(1.0, ), (-1.0, ), (2.0, ), (-2.0, )],
+    [(-2.0,)],
+)
+@pytest.mark.parametrize(
+    ("angle_x", ),
+    # [(0.0, ), (0.2, ), (-0.05, )],
+    [(-0.05,)]
+)
+@pytest.mark.parametrize(
+    ("angle_y", ),
+    [(0.0, ), (0.2, ), (-0.05, )],
+)
+@pytest.mark.parametrize(
+    ("angle_z", ),
+    [(0.0, ), (0.2, ), (-0.05, )],
+)
+@pytest.mark.parametrize(
+    ("offsets", ),
+    [
+        (None, ),
+        ((0.0, 0.0, 0.0),),
+    ],
+)
+def test_decimation(
+    tmp_path,
+    size,
+    block_size,
+    zoom_x,
+    zoom_y,
+    zoom_z,
+    angle_x,
+    angle_y,
+    angle_z,
+    offsets,
+):
     """Exercise decimation."""
 
     # Calculate the number of sub-blocks in each dimension
     num_blocks = [s // b for s, b in zip(size, block_size)]
 
     # Create the empty array
-    voxel_array = np.zeros(size, dtype=int)
+    voxel_array = np.zeros(size, dtype=np.uint16)
 
     # Fill the array with increasing values based on sub-block position
     current_block = 0
@@ -58,11 +105,23 @@ def test_decimation(tmp_path, size, block_size):
 
     fname = tmp_path / "test_img.nii.gz"
 
-    nb.Nifti1Image(voxel_array, None, None).to_filename(fname)
+    affine = np.eye(4)
+    affine[:3, :3] = (
+        nb.eulerangles.euler2mat(x=angle_x, y=angle_y, z=angle_z)
+        @ np.diag((zoom_x, zoom_y, zoom_z))
+        @ affine[:3, :3]
+    )
 
-    # Need to define test oracle. For now, just see if it doesn't smoke.
-    decimate(fname, factor=2, smooth=False, order=1)
+    if offsets is None:
+        affine[:3, 3] = -0.5 * nb.affines.apply_affine(affine, np.array(size) - 1)
 
-    # out.to_filename(tmp_path / "decimated.nii.gz")
+    test_image = nb.Nifti1Image(voxel_array.astype(np.uint16), affine, None)
+    test_image.header.set_data_dtype(np.uint16)
+    test_image.to_filename(fname)
+
+    # Need to define test oracle. For now, just see if it doesn't smoke.
+    out = decimate(fname, factor=2, smooth=False, order=1)
+    out.to_filename(tmp_path / "decimated.nii.gz")
 
-    # import pdb; pdb.set_trace()
+    out = downsample(fname, shape=(10, 10, 10), smooth=False, order=1)
+    out.to_filename(tmp_path / "downsampled.nii.gz")