mrid-python 0.1.3__py3-none-any.whl → 0.1.4__py3-none-any.whl

mrid/atlas/MNI152/__init__.py

@@ -62,6 +62,13 @@ def get_mni152(
 ):
     """Returns path to .nii.gz file of specified MNI-152 template.
 
+    The following templates are available:
+    - ``"2006 T1w symmetric"``
+    - ``"2009a T1w symmetric"``
+    - ``"2009a T2w symmetric"``
+    - ``"2009a T1w asymmetric"``
+    - ``"2009a T2w asymmetric"``
+
     Descriptions of templates are available here https://zenodo.org/records/15470657
     """
     filename = f"{type} {bool(skullstripped)}.nii.gz"

mrid/preprocessing/CTseg.py

@@ -0,0 +1,81 @@
+""""
+This requires CTseg docker image to be present in the system.
+
+If you haven't already, install docker to your OS following this https://docs.docker.com/engine/install/
+
+CTseg is available in the following repository: https://github.com/WCHN/CTseg
+
+Navigate to any folder, open the terminal in that folder and type in ``git clone https://github.com/WCHN/CTseg``. You might have to [install git](https://git-scm.com/install/) if you don't have it installed. This will create a new directory called ``CTseg`` in the folder and download the repository to it. Alternatively you can open the repository in your web browser, click on the green "Code" button near the top and click "Download ZIP", and unpack the archive to a folder named ``CTseg``.
+
+Then type in the command specified in "Docker" section of the read-me in https://github.com/WCHN/CTseg to build an image from the Dockerfile in this repository. I decided to not copy the command here in case it gets updated in the repository. This will download and build a docker image, usually to ``/var/lib/docker/`` on Linux, or inside the Docker Desktop virtual machine disk image on Windows. Note that this will download 3 GB of data, so it might take some time.
+
+After it is done, the CTseg functions from mrid can be used.
+"""
+import os
+import subprocess
+from pathlib import Path
+
+from ..loading import tositk
+
+def run_CTseg(
+    pth_ct: str | os.PathLike,
+    dir_out: str = "",
+    docker_image = "ubuntu:ctseg",
+) -> None:
+    """Runs ``CTseg`` command-line routine via ``subprocess.run``.
+
+    Args:
+        pth_ct (str | os.PathLike): path to a file which must be in a ``*.nii`` format.
+        dir_out (str, optional):
+            optional name of a directory that will be created next to ``path_ct`` nii file to save CTseg outputs to.
+            If empty, outputs are saved next to ``path_ct`` file. Defaults to ''.
+        docker_image (str, optional): name of the docker image that ``CTseg`` is installed in. Defaults to "ubuntu:ctseg".
+    """
+
+    # docker run --rm -it -v "/home/jj/data":/data ubuntu:ctseg function spm_CTseg '/data/CT.nii'
+    # better
+    # docker run --rm -it -v "/home/jj/data":/data ubuntu:ctseg eval "spm_CTseg('/data/CT.nii', '', true, true, true, true, 1.0)"
+    pth_ct = Path(pth_ct)
+
+    if dir_out != "":
+        if "/" in dir_out or "\\" in dir_out:
+            raise RuntimeError(
+                "dir_out should be name of directory that will be created next to `path_ct`. "
+                f"It can't be a path. Got '{dir_out}'"
+            )
+
+        dir_out = f"/data/{dir_out}"
+
+    command = [
+        "docker",
+        "run",
+
+        # --rm automatically removes the container's file system after the container exits. This is useful for running temporary containers.
+        "--rm",
+
+        # This is a combination of two flags, -i and -t:
+            # -i (interactive): Keeps the standard input (STDIN) open, allowing you to interact with the container.
+            # -t (tty): Allocates a pseudo-TTY, which makes the container behave like a normal terminal session. (doesn't work with subprocess)
+        #"-it",
+        "-i",
+
+        # -v used for mounting volumes, which allows you to connect a file path on your host machine to a path inside the container. This option requires a specific format: -v <host_path>:<container_path>.
+        "-v",
+        f"{os.path.normpath(pth_ct.parent)}:/data",
+
+        # docker image name
+        docker_image,
+
+        # evaluate matlab code
+        "eval",
+
+        # code to evaluate
+        f"spm_CTseg('/data/{pth_ct.name}', '{dir_out}', true, true, true, true, 1.0)",
+    ]
+
+    # run dcm2niix
+    subprocess.run(command, check=True)
+
+# this creates
+# wc01_1_00001_temp_CT_CTseg.nii
+# wc02_1_00001_temp_CT_CTseg

mrid/preprocessing/__init__.py

@@ -1,12 +1,12 @@
 from .bias_field_correction import n4_bias_field_correction
 from .cropping import crop_bg, crop_bg_D
-from .registration import resample_to, register, register_D, register_each, resize, downsample, Registration
-from .skullstripping import skullstrip, skullstrip_D, run_hd_bet, predict_brain_mask
+from .spatial import downsample, resample_to, resize
 
+# lib wrappers
+from . import hd_bet, CTseg, simple_elastix, synthstrip, mask
 __all__ = [
     "n4_bias_field_correction",
     "crop_bg", "crop_bg_D",
-    "resample_to", "register", "register_D", "register_each", "resize", "downsample",
-    "skullstrip", "skullstrip_D", "run_hd_bet", "predict_brain_mask"
-
-]
+    "resample_to", "resize", "downsample",
+    "hd_bet", "CTseg", "simple_elastix", "synthstrip", "mask",
+]

mrid/preprocessing/hd_bet.py

@@ -0,0 +1,242 @@
+import os
+import subprocess
+import tempfile
+from collections.abc import Mapping
+from typing import Literal
+
+import SimpleITK as sitk
+
+from ..loading.convert import ImageLike, tositk
+from ..utils.torch_utils import CUDA_IF_AVAILABLE
+from .simple_elastix import register, register_D
+from .mask import expand_binary_mask, apply_mask
+
+# hd_bet -h
+
+# -i INPUT, --input INPUT
+#                         input. Can be either a single file name or an input folder. If file: must be nifti (.nii.gz) and can only be 3D. No support for 4d images, use fslsplit to split 4d sequences into 3d images. If folder: all files ending with .nii.gz within that folder will be brain extracted.
+# -o OUTPUT, --output OUTPUT
+#                         output. Can be either a filename or a folder. If it does not exist, the folder will be created
+# -device DEVICE        used to set on which device the prediction will run. Can be 'cuda' (=GPU), 'cpu' or 'mps'. Default: cuda
+# --disable_tta         Set this flag to disable test time augmentation. This will make prediction faster at a slight decrease in prediction quality. Recommended for device cpu
+# --save_bet_mask       Set this flag to keep the bet masks. Otherwise they will be removed once HD_BET is done
+# --no_bet_image        Set this flag to disable generating the skull stripped/brain extracted image. Only makes sense if you also set --save_bet_mask
+# --verbose             Talk to me.
+
+def run_hd_bet(
+    input: str | os.PathLike,
+    output: str | os.PathLike,
+    device: Literal['cpu', 'cuda', 'mps'] = CUDA_IF_AVAILABLE,
+    disable_tta: bool = False,
+    save_bet_mask: bool = True,
+    no_bet_image: bool = False,
+    verbose: bool = False,
+) -> None:
+    """Runs HD-BET command-line routine via ``subprocess.run``.
+
+    Args:
+        input (str | os.PathLike):
+            input. Can be either a single file name or an input folder.
+            If file: must be nifti (.nii.gz) and can only be 3D.
+            No support for 4d images, use fslsplit to split 4d sequences into 3d images.
+            If folder: all files ending with .nii.gz within that folder will be brain extracted.
+        output (str | os.PathLike):
+            output. Can be either a filename or a folder. If it does not exist, the folder will be created
+        device (str, optional):
+            used to set on which device the prediction will run.
+            Can be 'cuda' (=GPU), 'cpu' or 'mps'. Default: cuda. Defaults to CUDA_IF_AVAILABLE.
+        disable_tta (bool, optional):
+            Set this flag to disable test time augmentation.
+            This will make prediction faster at a slight decrease in prediction quality.
+            Recommended for device cpu. Defaults to False.
+        save_bet_mask (bool, optional):
+            Set this flag to keep the bet masks.
+            Otherwise they will be removed once HD_BET is done. Defaults to True.
+        no_bet_image (bool, optional):
+            Set this flag to disable generating the skull stripped/brain extracted image.
+            Only makes sense if you also set --save_bet_mask. Defaults to False.
+        verbose (bool, optional): Talk to me. Defaults to False.
+    """
+
+    command = [
+        "hd-bet",
+        "-i", os.path.normpath(input),
+        "-o", os.path.normpath(output),
+        "-device", device,
+    ]
+    if disable_tta: command.append("--disable_tta")
+    if save_bet_mask: command.append("--save_bet_mask")
+    if no_bet_image: command.append("--no_bet_image")
+    if verbose: command.append("--verbose")
+
+    # run dcm2niix
+    subprocess.run(command, check=True)
+
+
+def predict_brain_mask(
+    input: ImageLike,
+    register_to_mni152: Literal["T1", "T2"] | None = None,
+    device: Literal["cpu", "cuda", "mps"] = CUDA_IF_AVAILABLE,
+    disable_tta: bool = False,
+    verbose: bool = False,
+) -> sitk.Image:
+    """Returns brain mask of ``input`` predicted by HD-BET.
+
+    Args:
+        input (ImageLike): input to skullstrip. Recommended T1-w, postcontrast T1-w, T2-w or FLAIR sequence in MNI152 space.
+        register_to_mni152 (str | None, optional):
+            Modality of MNI152 template to pre-register ``input`` to. Should be ``"T1"``, ``"T2"`` or ``None``.
+            if specified, ``input`` will be registered to specified MNI152 template,
+            then after prediction the brain mask registered back to original ``input``.
+            Note that HD-BET expects images to be in MNI152 space. Defaults to None.
+        device (str, optional):
+            used to set on which device the prediction will run. Can be 'cuda' (=GPU), 'cpu' or 'mps'.
+            Defaults to CUDA_IF_AVAILABLE.
+        disable_tta (bool, optional):
+            Set this flag to disable test time augmentation.
+            This will make prediction faster at a slight decrease in prediction quality.
+            Recommended for device cpu. Defaults to False.
+        verbose (bool, optional): Talk to me. Defaults to False.
+    """
+    input = tositk(input)
+
+    # ---------------------------- register to mni152 ---------------------------- #
+    if register_to_mni152 is not None:
+        from ..atlas.MNI152 import get_mni152
+        mni152 = get_mni152(f"2009a {register_to_mni152}w asymmetric", skullstripped=False) # type:ignore
+        input_mni = register(input, mni152)
+
+    else:
+        input_mni = input
+
+    # ---------------------------- predict brain mask ---------------------------- #
+    with tempfile.TemporaryDirectory() as tmpdir:
+        sitk.WriteImage(input_mni, os.path.join(tmpdir, "input.nii.gz"))
+
+        run_hd_bet(
+            input = os.path.join(tmpdir, "input.nii.gz"),
+            output = os.path.join(tmpdir, "output.nii.gz"),
+            device=device, disable_tta=disable_tta, save_bet_mask=True, verbose=verbose,
+        )
+
+        brain_mask_mni = tositk(os.path.join(tmpdir, "output_bet.nii.gz"))
+
+    # ------------------------- unregister mask if needed ------------------------ #
+    if register_to_mni152 is not None:
+        study_mni = dict(image=input_mni, seg_brain=brain_mask_mni)
+        study = register_D(study_mni, key="image", to=input)
+        brain_mask = study["seg_brain"]
+
+    else:
+        brain_mask = brain_mask_mni
+
+    return brain_mask
+
+def skullstrip(
+    input: ImageLike,
+    register_to_mni152: Literal["T1", "T2"] | None = None,
+    device: Literal["cpu", "cuda", "mps"] = CUDA_IF_AVAILABLE,
+    disable_tta: bool = False,
+    verbose: bool = False,
+
+    expand: int = 0,
+) -> sitk.Image:
+    """Skullstrips ``input`` using HD-BET.
+
+    Args:
+        input (ImageLike): input to skullstrip. Recommended T1-w, postcontrast T1-w, T2-w or FLAIR sequence in MNI152 space.
+        register_to_mni152 (str | None, optional):
+            Modality of MNI152 template to pre-register ``input`` to. Should be ``"T1"``, ``"T2"`` or ``None``.
+            if specified, ``input`` will be registered to specified MNI152 template,
+            then after prediction the brain mask registered back to original ``input``.
+            Note that HD-BET expects images to be in MNI152 space. Defaults to None.
+        device (str, optional):
+            used to set on which device the prediction will run. Can be 'cuda' (=GPU), 'cpu' or 'mps'.
+            Defaults to CUDA_IF_AVAILABLE.
+        disable_tta (bool, optional):
+            Set this flag to disable test time augmentation. This will make prediction faster
+            at a slight decrease in prediction quality. Recommended for device cpu. Defaults to False.
+        verbose (bool, optional): Talk to me. Defaults to False.
+        expand (int, optional):
+            Positive values expand brain mask by this many pixels, meaning inner parts of the skull will be included;
+            Negative values dilate brain mask by this many pixels, meaning outer parts of the brain will be excluded.
+
+    """
+    input = tositk(input)
+    mask = predict_brain_mask(input=input, register_to_mni152=register_to_mni152,
+                                  device=device, disable_tta=disable_tta, verbose=verbose)
+
+    if expand != 0:
+        mask = expand_binary_mask(mask, expand=expand)
+
+    return apply_mask(input, mask)
+
+
+def skullstrip_D(
+    images: Mapping[str, ImageLike],
+    key: str,
+    register_to_mni152: Literal["T1", "T2"] | None = None,
+    device: Literal["cpu", "cuda", "mps"] = CUDA_IF_AVAILABLE,
+    disable_tta: bool = False,
+    verbose: bool = False,
+
+    expand: int = 0,
+
+    include_mask: bool = False,
+    keep_original: bool = False,
+) -> dict[str, sitk.Image]:
+    """Predicts brain mask of ``images[key]`` using HD-BET, then uses this mask to skull strip all values in ``images``.
+
+    Args:
+        images (Mapping[str, ImageLike]): dictionary of images that align with each other.
+        key (str): key of the image to pass to HD-BET for brain mask prediction.
+        register_to_mni152 (str | None, optional):
+            Modality of MNI152 template to pre-register ``input`` to. Should be ``"T1"``, ``"T2"`` or ``None``.
+            if specified, ``input`` will be registered to specified MNI152 template,
+            then after prediction the brain mask registered back to original ``input``.
+            Note that HD-BET expects images to be in MNI152 space. Defaults to None.
+        device (str, optional):
+            used to set on which device the prediction will run. Can be 'cuda' (=GPU), 'cpu' or 'mps'.
+            Defaults to CUDA_IF_AVAILABLE.
+        disable_tta (bool, optional):
+            Set this flag to disable test time augmentation. This will make prediction faster
+            at a slight decrease in prediction quality. Recommended for device cpu. Defaults to False.
+        verbose (bool, optional): Talk to me. Defaults to False.
+        expand (int, optional):
+            Positive values expand brain mask by this many pixels, meaning inner parts of the skull will be included;
+            Negative values dilate brain mask by this many pixels, meaning outer parts of the brain will be excluded.
+        include_mask (bool, optional):
+            if True, adds ``"seg_seg_hd_bet"`` with brain mask predicted by HD-BET to returned dictionary.
+            This adds brain mask BEFORE expanding/dilating if ``expand`` argument is specified.
+        keep_original (bool, Optional):
+            if True, skull-stripped images are added to the dictionary
+            with ``"_hd_bet" ``postfix, rather than replacing.
+    """
+    images = {k: tositk(v) for k,v in images.items()}
+
+    mask = predict_brain_mask(input=images[key], register_to_mni152=register_to_mni152,
+                          device=device, disable_tta=disable_tta, verbose=verbose)
+
+    skullstripped = {}
+
+    # include mask before expanding
+    if include_mask:
+        mask_sitk = tositk(mask)
+        mask_sitk.CopyInformation(images[key])
+        skullstripped["seg_hd_bet"] = mask_sitk
+
+    # expand
+    if expand != 0:
+        mask = expand_binary_mask(mask, expand=expand)
+
+    # apply mask
+    for k,v in images.items():
+        skullstripped[k] = apply_mask(v, mask)
+
+    # optionally add with skullstripped postfix
+    if keep_original:
+        skullstripped = {f"{k}_hd_bet": v for k,v in skullstripped}
+        skullstripped.update(images.copy())
+
+    return skullstripped
+

mrid/preprocessing/mask.py

@@ -0,0 +1,42 @@
+import SimpleITK as sitk
+import numpy as np
+import SimpleITK as sitk
+
+from ..loading.convert import ImageLike, tositk, tonumpy
+
+
+def expand_binary_mask(binary_mask: ImageLike, expand: int) -> sitk.Image:
+    """Expand or dilate a binary mask.
+
+    Args:
+        binary_mask (ImageLike): mask
+        expand (int, optional):
+            Positive values expand the mask by this many pixels;
+            Negative values dilate the mask by this many pixels.
+    """
+    binary_mask = tositk(binary_mask)
+    if expand > 0:
+        inverted_mask = 1 - binary_mask
+        return 1 - sitk.BinaryDilate(inverted_mask, (expand, expand, expand))
+
+    if expand < 0:
+        return sitk.BinaryDilate(binary_mask, (-expand, -expand, -expand))
+
+    return binary_mask
+
+def apply_mask(image: ImageLike, mask: ImageLike) -> sitk.Image:
+    """Applies ``mask`` to ``image``, that is all values where ``mask > 0`` are kept.
+
+    This function sets all values outside of the mask to smallest value within the mask."""
+    image = tositk(image)
+    mask = tositk(mask)
+
+    mask = sitk.Cast(mask, image.GetPixelID())
+
+    image_np = sitk.GetArrayFromImage(image)
+    mask_np = (tonumpy(mask) > 0).astype(np.bool)
+    image_ma = np.ma.masked_array(image_np, ~mask_np)
+
+    image_applied = tositk(image_ma.filled(image_ma.min()))
+    image_applied.CopyInformation(image)
+    return image_applied

mrid/preprocessing/{registration.py → simple_elastix.py}

@@ -1,22 +1,12 @@
 from collections.abc import Mapping, Sequence
 from typing import TYPE_CHECKING, Any
 
+import os
 import numpy as np
 import SimpleITK as sitk
 
 from ..loading.convert import tositk, ImageLike
 
-def resample_to(input: ImageLike, to: ImageLike, interpolation=sitk.sitkNearestNeighbor) -> sitk.Image:
-    """Resample ``input`` to ``reference``.
-
-    Resampling uses spatial information embedded in the sitk.Image - size, origin, spacing and direction.
-
-    Note that this information is only available when certain imaging formats are loaded, such as DICOM and NIfTI.
-
-    ``input`` is transformed in such a way that those attributes will match ``reference``.
-    """
-    return sitk.Resample(tositk(input), tositk(to), sitk.Transform(), interpolation)
-
 
 def _default_pmap():
     """Default parameter maps for registration"""
@@ -29,8 +19,8 @@ def _default_pmap():
     pmap.append(sitk.GetDefaultParameterMap("affine"))
     return pmap
 
-class Registration:
-    """Class for image registration.
+class SimpleElastix:
+    """Class for image registration via SimpleElastix.
 
     Args:
         pmap (Any, optional): parameter map, if None, uses default parameter map. Defaults to None.
@@ -50,7 +40,7 @@ class Registration:
 
         self._moving = None
         self._transformed = None
-        self.inverse: "Registration | None" = None
+        self.inverse: "SimpleElastix | None" = None
 
     def find_transform(self, input: ImageLike, to: ImageLike) -> sitk.Image:
         """Find a transform that transforms ``input`` to ``to`` and save it to this ``Registration`` object.
@@ -62,6 +52,7 @@ class Registration:
         """
         if self._transformed is not None:
             raise RuntimeError("`find_transform` has already been called on this Registration object.")
+
         self._moving = tositk(input)
         to = tositk(to)
 
@@ -88,6 +79,9 @@ class Registration:
         if self._transformed is None:
             raise RuntimeError("First find transform parameters using `find_transform` method.")
 
+        input = tositk(input)
+        input.CopyInformation(self._moving)
+
         transform = sitk.TransformixImageFilter()
         tmap = self.elastix.GetTransformParameterMap()
         if use_nearest_interpolation:
@@ -118,7 +112,7 @@ class Registration:
             inverse_pmap = self.elastix.GetParameterMap() # this returns a copy
             # for p in inverse_pmap:
             #     p["Metric"] = "MeanSquaredDifference" # not implemented
-            self.inverse = Registration(pmap=inverse_pmap, log_to_console=self.log_to_console)
+            self.inverse = SimpleElastix(pmap=inverse_pmap, log_to_console=self.log_to_console)
             self.inverse.find_transform(
                 input=self._transformed,
                 to=self._moving
@@ -138,7 +132,7 @@ def register(input: ImageLike, to: ImageLike, pmap: Any = None, log_to_console=F
     and install https://pypi.org/project/SimpleITK-SimpleElastix/, don't worry, it's
     the same as SimpleITK but it additionally includes SimpleElastix.
     """
-    reg = Registration(pmap=pmap, log_to_console=log_to_console)
+    reg = SimpleElastix(pmap=pmap, log_to_console=log_to_console)
     return reg.find_transform(input=input, to=to)
 
 
@@ -155,7 +149,7 @@ def register_D(
     Make sure segmentation with hard edges is under a key that starts with ``"seg"``,
     it will use nearest neighbour interpolation, otherwise it will mess up the edges.
     """
-    reg = Registration(pmap=pmap, log_to_console=log_to_console)
+    reg = SimpleElastix(pmap=pmap, log_to_console=log_to_console)
     registered = {key: reg.find_transform(images[key], to)}
 
     # process segs last because it sets resample interpolator to nearest
@@ -169,18 +163,23 @@ def register_D(
 def register_each(
     images: Mapping[str, ImageLike],
     key: str,
-    to: ImageLike,
+    to: "ImageLike | None" = None,
     pmap: Any = None,
     log_to_console=False,
 ) -> dict[str, sitk.Image]:
-    """Register ``images[key]`` to ``reference``, then register all other values in ``images`` to registered ``images[key]``.
+    """Registers all other images to ``images[key]``.
+    If ``to`` is specified, register ``images[key]`` to ``to`` beforehand.
+    Uses SimpleElastix.
 
     Use this when you have multiple modalities that do not align."""
     images = {k: tositk(v) for k,v in images.items()}
-    to = tositk(to)
 
     input = images[key]
-    input_reg = register(input=input, to=to, pmap=pmap, log_to_console=log_to_console)
+    if to is not None:
+        to = tositk(to)
+        input_reg = register(input=input, to=to, pmap=pmap, log_to_console=log_to_console)
+    else:
+        input_reg = input
 
     registered = {key: input_reg}
     for k,v in images.items():
@@ -188,64 +187,3 @@ def register_each(
             registered[k] = register(input=v, to=input_reg, pmap=pmap, log_to_console=log_to_console)
 
     return registered
-
-def resize(img: ImageLike, new_size: Sequence[int], interpolator=sitk.sitkLinear) -> sitk.Image:
-    """Resize ``sitk.Image`` to ``new_size``. Retains correct spatial information.
-    source: https://gist.github.com/lixinqi98/1bbd3596492f20b776fed2778f7cd48c"""
-    img = tositk(img)
-    new_size = list(reversed(new_size))
-
-    # img = sitk.ReadImage(img)
-    dimension = img.GetDimension()
-
-    # Physical image size corresponds to the largest physical size in the training set, or any other arbitrary size.
-    reference_physical_size = np.zeros(dimension)
-
-    reference_physical_size[:] = [(sz - 1) * spc if sz * spc > mx else mx for sz, spc, mx in
-                                  zip(img.GetSize(), img.GetSpacing(), reference_physical_size)]
-
-    # Create the reference image with a zero origin, identity direction cosine matrix and dimension
-    reference_origin = np.zeros(dimension)
-    reference_direction = np.identity(dimension).flatten()
-    reference_size = new_size
-    reference_spacing = [phys_sz / (sz - 1) for sz, phys_sz in zip(reference_size, reference_physical_size)]
-
-    reference_image = sitk.Image(reference_size, img.GetPixelIDValue())
-    reference_image.SetOrigin(reference_origin)
-    reference_image.SetSpacing(reference_spacing)
-    reference_image.SetDirection(reference_direction)
-
-    # Always use the TransformContinuousIndexToPhysicalPoint to compute an indexed point's physical coordinates as
-    # this takes into account size, spacing and direction cosines. For the vast majority of images the direction
-    # cosines are the identity matrix, but when this isn't the case simply multiplying the central index by the
-    # spacing will not yield the correct coordinates resulting in a long debugging session.
-    reference_center = np.array(
-        reference_image.TransformContinuousIndexToPhysicalPoint(np.array(reference_image.GetSize()) / 2.0))
-
-    # Transform which maps from the reference_image to the current img with the translation mapping the image
-    # origins to each other.
-    transform = sitk.AffineTransform(dimension)
-    transform.SetMatrix(img.GetDirection())
-    transform.SetTranslation(np.array(img.GetOrigin()) - reference_origin)
-    # Modify the transformation to align the centers of the original and reference image instead of their origins.
-    centering_transform = sitk.TranslationTransform(dimension)
-    img_center = np.array(img.TransformContinuousIndexToPhysicalPoint(np.array(img.GetSize()) / 2.0))
-    centering_transform.SetOffset(np.array(transform.GetInverse().TransformPoint(img_center) - reference_center))
-
-    # centered_transform = sitk.Transform(transform)
-    # centered_transform.AddTransform(centering_transform)
-
-    centered_transform = sitk.CompositeTransform([transform, centering_transform])
-
-    # Using the linear interpolator as these are intensity images, if there is a need to resample a ground truth
-    # segmentation then the segmentation image should be resampled using the NearestNeighbor interpolator so that
-    # no new labels are introduced.
-
-    return sitk.Resample(img, reference_image, centered_transform, interpolator, 0.0)
-
-def downsample(image:ImageLike, factor:float, dims: Sequence[int] | None, interpolator=sitk.sitkLinear) -> sitk.Image:
-    """factor = 2 for 2x downsampling"""
-    image = tositk(image)
-    size = sitk.GetArrayFromImage(image).shape
-    size = [round(s/factor) if (dims is None or i in dims) else s for i,s in enumerate(size)]
-    return resize(image, size, interpolator=interpolator)

mrid/preprocessing/spatial.py

@@ -0,0 +1,81 @@
+from collections.abc import Mapping, Sequence
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+import SimpleITK as sitk
+
+from ..loading.convert import tositk, ImageLike
+
+
+def resample_to(input: ImageLike, to: ImageLike, interpolation=sitk.sitkNearestNeighbor) -> sitk.Image:
+    """Resample ``input`` to ``reference``.
+
+    Resampling uses spatial information embedded in the sitk.Image - size, origin, spacing and direction.
+
+    Note that this information is only available when certain imaging formats are loaded, such as DICOM and NIfTI.
+
+    ``input`` is transformed in such a way that those attributes will match ``reference``.
+    """
+    return sitk.Resample(tositk(input), tositk(to), sitk.Transform(), interpolation)
+
+
+def resize(img: ImageLike, new_size: Sequence[int], interpolator=sitk.sitkLinear) -> sitk.Image:
+    """Resize ``sitk.Image`` to ``new_size``. Retains correct spatial information.
+    source: https://gist.github.com/lixinqi98/1bbd3596492f20b776fed2778f7cd48c"""
+    img = tositk(img)
+    new_size = list(reversed(new_size))
+
+    # img = sitk.ReadImage(img)
+    dimension = img.GetDimension()
+
+    # Physical image size corresponds to the largest physical size in the training set, or any other arbitrary size.
+    reference_physical_size = np.zeros(dimension)
+
+    reference_physical_size[:] = [(sz - 1) * spc if sz * spc > mx else mx for sz, spc, mx in
+                                  zip(img.GetSize(), img.GetSpacing(), reference_physical_size)]
+
+    # Create the reference image with a zero origin, identity direction cosine matrix and dimension
+    reference_origin = np.zeros(dimension)
+    reference_direction = np.identity(dimension).flatten()
+    reference_size = new_size
+    reference_spacing = [phys_sz / (sz - 1) for sz, phys_sz in zip(reference_size, reference_physical_size)]
+
+    reference_image = sitk.Image(reference_size, img.GetPixelIDValue())
+    reference_image.SetOrigin(reference_origin)
+    reference_image.SetSpacing(reference_spacing)
+    reference_image.SetDirection(reference_direction)
+
+    # Always use the TransformContinuousIndexToPhysicalPoint to compute an indexed point's physical coordinates as
+    # this takes into account size, spacing and direction cosines. For the vast majority of images the direction
+    # cosines are the identity matrix, but when this isn't the case simply multiplying the central index by the
+    # spacing will not yield the correct coordinates resulting in a long debugging session.
+    reference_center = np.array(
+        reference_image.TransformContinuousIndexToPhysicalPoint(np.array(reference_image.GetSize()) / 2.0))
+
+    # Transform which maps from the reference_image to the current img with the translation mapping the image
+    # origins to each other.
+    transform = sitk.AffineTransform(dimension)
+    transform.SetMatrix(img.GetDirection())
+    transform.SetTranslation(np.array(img.GetOrigin()) - reference_origin)
+    # Modify the transformation to align the centers of the original and reference image instead of their origins.
+    centering_transform = sitk.TranslationTransform(dimension)
+    img_center = np.array(img.TransformContinuousIndexToPhysicalPoint(np.array(img.GetSize()) / 2.0))
+    centering_transform.SetOffset(np.array(transform.GetInverse().TransformPoint(img_center) - reference_center))
+
+    # centered_transform = sitk.Transform(transform)
+    # centered_transform.AddTransform(centering_transform)
+
+    centered_transform = sitk.CompositeTransform([transform, centering_transform])
+
+    # Using the linear interpolator as these are intensity images, if there is a need to resample a ground truth
+    # segmentation then the segmentation image should be resampled using the NearestNeighbor interpolator so that
+    # no new labels are introduced.
+
+    return sitk.Resample(img, reference_image, centered_transform, interpolator, 0.0)
+
+def downsample(image:ImageLike, factor:float, dims: Sequence[int] | None, interpolator=sitk.sitkLinear) -> sitk.Image:
+    """factor = 2 for 2x downsampling"""
+    image = tositk(image)
+    size = sitk.GetArrayFromImage(image).shape
+    size = [round(s/factor) if (dims is None or i in dims) else s for i,s in enumerate(size)]
+    return resize(image, size, interpolator=interpolator)