mrid-python 0.1.5__tar.gz → 0.1.7__tar.gz

{mrid_python-0.1.5 → mrid_python-0.1.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mrid-python
-Version: 0.1.5
+Version: 0.1.7
 Summary: Tools for working with 3D medical images and segmentations - registration, brain skull-stripping, etc.
 Author-email: Ivan Nikishev <nkshv2@gmail.com>
 Project-URL: Homepage, https://github.com/inikishev/mrid

{mrid_python-0.1.5 → mrid_python-0.1.7}/mrid/atlas/MNI152/__init__.py

@@ -60,7 +60,7 @@ def get_mni152(
     ],
     skullstripped: bool = False,
 ):
-    """Returns path to .nii.gz file of specified MNI-152 template.
+    """Returns path to .nii.gz file of specified MNI-152 template. All templates are with skull.
 
     The following templates are available:
     - ``"2006 T1w symmetric"``

{mrid_python-0.1.5 → mrid_python-0.1.7}/mrid/preprocessing/CTseg.py

@@ -73,7 +73,7 @@ def run_CTseg(
         f"spm_CTseg('/data/{pth_ct.name}', '{dir_out}', true, true, true, true, 1.0)",
     ]
 
-    # run dcm2niix
+    # run
     subprocess.run(command, check=True)
 
 # this creates

{mrid_python-0.1.5 → mrid_python-0.1.7}/mrid/preprocessing/__init__.py

@@ -1,9 +1,9 @@
 from .bias_field_correction import n4_bias_field_correction
-from .cropping import crop_bg, crop_bg_D
+from .cropping import crop_bg, crop_bg_D, center_crop_or_pad
 from .spatial import downsample, resample_to, resize
 
 # lib wrappers
-from . import hd_bet, CTseg, simple_elastix, synthstrip, mask
+from . import hd_bet, CTseg, simple_elastix, synthstrip, mask, haca3
 __all__ = [
     "n4_bias_field_correction",
     "crop_bg", "crop_bg_D",

{mrid_python-0.1.5 → mrid_python-0.1.7}/mrid/preprocessing/cropping.py

@@ -1,4 +1,4 @@
-from collections.abc import Mapping
+from collections.abc import Mapping, Sequence
 from typing import Any
 import SimpleITK as sitk
 
@@ -34,3 +34,39 @@ def crop_bg_D(images: Mapping[str, ImageLike], key: str) -> dict[str, sitk.Image
     ret = {k: sitk.RegionOfInterest(v, bbox[int(len(bbox) / 2) :],  bbox[0 : int(len(bbox) / 2)]) for k,v in images.items()}
     return ret
 
+def center_crop_or_pad(image, size: Sequence[int]) -> "sitk.Image":#[192, 224, 192]
+    """Crops or pads image from the center to ``size``."""
+    image = tositk(image)
+    current_size = list(image.GetSize())
+
+    low_pad = []
+    high_pad = []
+    low_crop = []
+    high_crop = []
+
+    for cur, tar in zip(current_size, size):
+        diff = tar - cur
+        if diff >= 0:
+            low = diff // 2
+            high = diff - low
+            low_pad.append(low)
+            high_pad.append(high)
+            low_crop.append(0)
+            high_crop.append(0)
+        else:
+            diff = abs(diff)
+            low = diff // 2
+            high = diff - low
+            low_pad.append(0)
+            high_pad.append(0)
+            low_crop.append(low)
+            high_crop.append(high)
+
+    image = sitk.ConstantPad(image, low_pad, high_pad, 0)
+    image = sitk.Crop(image, low_crop, high_crop)
+
+    # Verify size
+    if tuple(image.GetSize()) == tuple(size):
+        return image
+
+    raise RuntimeError(f"Final size is {image.GetSize()} instead of {size}, crop_or_pad failed for some reason.")

mrid_python-0.1.7/mrid/preprocessing/haca3.py

@@ -0,0 +1,242 @@
+""""
+This requires [HACA3](https://github.com/lianruizuo/haca3).
+
+HACA3 can be installed from source or through singularity image. Mrid currently only supports it being installed in a separate environment (because thats how I currently need to use it), let me know if you need other installation methods supported.
+
+You also need to download HACA3 weights harmonization.pt and fusion model weights fusion.pt from the ``4. Usage: Inference`` section in HACA3 github readme.
+"""
+import os
+import shlex
+import subprocess
+import tempfile
+from collections.abc import Sequence
+from pathlib import Path
+from typing import Literal
+
+import SimpleITK as sitk
+
+from ..loading import ImageLike, tositk
+from .cropping import center_crop_or_pad
+from .simple_elastix import register, register_D
+
+# - ```--in-path```: file path to input source image. Multiple ```--in-path``` may be provided if there are multiple
+# source images. See the above example for more details.
+# - ```--target-image```: file path to target image. HACA3 will match the contrast of source images to this target image.
+# - ```--target-theta```: In HACA3, ```theta```
+# is a two-dimensional representation of image contrast. Target image contrast can be directly specified by providing
+# a ```theta``` value, e.g., ```--target-theta 0.5 0.5```. Note: either ```--target-image``` or ```--target-image``` must
+# be provided during inference. If both are provided, only ```--target-theta``` will be used.
+# - ```--norm-val```: normalization value.
+# - ```--out-path```: file path to harmonized image.
+# - ```--harmonization-model```: pretrained HACA3 weights. Pretrained model weights on IXI, OASIS and HCP data can
+# be downloaded [here](https://iacl.ece.jhu.edu/~lianrui/haca3/harmonization_public.pt).
+# - ```--fusion-model```: pretrained fusion model weights. HACA3 uses a 3D convolutional network to combine multi-orientation
+# 2D slices into a single 3D volume. Pretrained fusion model can be downloaded [here](https://iacl.ece.jhu.edu/~lianrui/haca3/fusion.pt).
+# - ```--save-intermediate```: if specified, intermediate results will be saved. Default: ```False```. Action: ```store_true```.
+# - ```--intermediate-out-dir```: directory to save intermediate results.
+# - ```--gpu-id```: integer number specifies which GPU to run HACA3.
+# - ```--num-batches```: During inference, HACA3 takes entire 3D MRI volumes as input. This may cause a considerable amount
+# GPU memory. For reduced GPU memory consumption, source images maybe divided into smaller batches.
+# However, this may slightly increase the inference time.
+
+
+def run_HACA3(
+    conda_path: str | os.PathLike,
+    env_name: str,
+    harmonization_model: str | os.PathLike,
+    fusion_model: str | os.PathLike,
+    in_path: str | os.PathLike | Sequence[str | os.PathLike],
+    out_path: str | os.PathLike,
+    target_image: str | os.PathLike | None,
+    target_theta: tuple[float,float] | None,
+    norm_val: float | None = None,
+    intermediate_out_dir: str | os.PathLike | None = None,
+    gpu_id: int | None = None,
+    num_batches: int | None = None,
+) -> None:
+    """Runs ``HACA3`` command-line routine via ``subprocess.run``.
+
+    Args:
+        conda_path: path to ``minconda3`` directory.
+        env_name: name of the conda env where HACA3 is installed.
+        harmonization_model: pretrained HACA3 weights. Pretrained model weights on IXI, OASIS and HCP
+            data can be downloaded [here](https://iacl.ece.jhu.edu/~lianrui/haca3/harmonization_public.pt).
+        fusion_model: pretrained fusion model weights. HACA3 uses a 3D convolutional network to
+            combine multi-orientation 2D slices into a single 3D volume. Pretrained fusion model
+            can be downloaded [here](https://iacl.ece.jhu.edu/~lianrui/haca3/fusion.pt).
+        in_path: file path to input source image. Multiple paths may be provided if there are
+            multiple source images (different modalities). Note that modalities must be in
+            MNI space (1mm isotropic resolution). HACA3 assumes a spatial dimension of 192x224x192.
+        out_path: file path to harmonized image.
+        target_image: file path to target image. HACA3 will match the contrast of
+            source images to this target image.
+        target_theta: In HACA3, ```theta``` is a two-dimensional representation of image contrast.
+            Target image contrast can be directly specified by providing a ```theta``` value, e.g.,
+            ```target_theta = (0.5, 0.5)```. Note: either ```target_image``` or ```target_theta```
+            must be provided during inference.
+        norm_val: normalization value. Defaults to None.
+        intermediate_out_dir: directory to save intermediate results. Defaults to None.
+        gpu_id: integer number specifies which GPU to run HACA3. Defaults to None.
+        num_batches: During inference, HACA3 takes entire 3D MRI volumes as input.
+            This may cause a considerable amount GPU memory. For reduced GPU memory consumption,
+            source images maybe divided into smaller batches. However, this may slightly
+            increase the inference time. Defaults to None.
+
+    """
+    # validate env name
+    env_name = str(env_name)
+    if not all((c.isalnum() or c in ("-_")) for c in env_name):
+        raise RuntimeError(f"env_name contains invalid characters: {env_name}")
+
+    haca3_command = ['haca3-test']
+
+    if isinstance(in_path, (str, os.PathLike)):
+        in_path = [in_path]
+
+    if target_theta is not None:
+        if not isinstance(target_theta, tuple):
+            raise RuntimeError(f"target_theta must be tuple of two float values or None, got {type(target_theta)}")
+        if not len(target_theta) == 2:
+            raise RuntimeError(f"target_theta must be a length 2 tuple, got length {len(target_theta)}")
+        if not all(isinstance(t, (int,float)) for t in target_theta):
+            raise RuntimeError(
+                f"target_theta must be tuple of two float values, got tuple({tuple(type(v) for v in target_theta)})")
+
+    for f in in_path:
+        haca3_command.append(f'--in-path "{os.path.normpath(f)}"',)
+
+    haca3_command.extend([
+        f'--out-path "{os.path.normpath(out_path)}"',
+        f'--harmonization-model "{os.path.normpath(harmonization_model)}"',
+        f'--fusion-model "{os.path.normpath(fusion_model)}"',
+    ])
+
+    if target_image is not None: haca3_command.append(f'--target-image "{os.path.normpath(target_image)}"')
+    if target_theta is not None: haca3_command.append(f'--target-theta {float(target_theta[0])} {float(target_theta[1])}')
+    if norm_val is not None: haca3_command.append(f'--norm-val {float(norm_val)}')
+    if intermediate_out_dir is not None:
+        haca3_command.append('--save-intermediate')
+        haca3_command.append(f'--intermediate-out-dir "{os.path.normpath(intermediate_out_dir)}"')
+    if gpu_id is not None: haca3_command.append(f'--gpu-id {int(gpu_id)}')
+    if num_batches is not None: haca3_command.append(f'--num-batches {int(num_batches)}')
+
+    # print(f". {conda_path}/etc/profile.d/conda.sh && conda activate {env_name} && {' '.join(haca3_command)}")
+
+    # so
+    # shlex.split doesn't work on ., and conda run doesn't work for whatever reason, so we have to do this
+    # all args are explicitly validated so should be ok
+    command = f". {conda_path}/etc/profile.d/conda.sh && conda activate {env_name} && {' '.join(haca3_command)}"
+
+    # run
+    subprocess.run(command, shell=True, check=True)
+
+def harmonize(
+    conda_path: str | os.PathLike,
+    env_name: str,
+    harmonization_model: str | os.PathLike,
+    fusion_model: str | os.PathLike,
+    inputs: "ImageLike | Sequence[ImageLike]",
+    target_image: "ImageLike | None" = None,
+    target_theta: tuple[float,float] | None = None,
+    norm_val: float | None = None,
+    intermediate_out_dir: str | os.PathLike | None = None,
+    gpu_id: int | None = None,
+    num_batches: int | None = None,
+) -> sitk.Image:
+    """Harmonizes ``inputs`` using HACA3.
+
+    Important: Some preprocessing steps are needed before running HACA3:
+
+    - Inhomogeneity correction
+    - Super-resolution for 2D acquired scans. This step is optional, but recommended for optimal performance. See [SMORE](https://github.com/volcanofly/SMORE-Super-resolution-for-3D-medical-images-MRI) for more details.
+    - Registration to MNI space (1mm isotropic resolution). HACA3 assumes a spatial dimension of 192x224x192.
+
+    You can inhomogeneity correction via ``mrid.n4_bias_field_correction(...)``,
+    register to MNI152 via ``mrid.simple_elastix.register(...)``, and center-pad to required dimension iva
+    ``mrid.center_crop_or_pad(image, [192, 224, 192])``.
+
+    Args:
+        conda_path: path to ``minconda3`` directory.
+        env_name: name of the conda env where HACA3 is installed.
+        harmonization_model: pretrained HACA3 weights. Pretrained model weights on IXI, OASIS and HCP
+            data can be downloaded [here](https://iacl.ece.jhu.edu/~lianrui/haca3/harmonization_public.pt).
+        fusion_model: pretrained fusion model weights. HACA3 uses a 3D convolutional network to
+            combine multi-orientation 2D slices into a single 3D volume. Pretrained fusion model
+            can be downloaded [here](https://iacl.ece.jhu.edu/~lianrui/haca3/fusion.pt).
+        inputs: input source images. Multiple images may be provided if there are
+            multiple source images (different modalities). Note that modalities must be in
+            MNI space (1mm isotropic resolution). HACA3 assumes a spatial dimension of 192x224x192.
+        target_image: target image. HACA3 will match the contrast of
+            source images to this target image. One of ``target_image`` or ``target_theta`` must be set.
+        target_theta: In HACA3, ```theta``` is a two-dimensional representation of image contrast.
+            Target image contrast can be directly specified by providing a ```theta``` value, e.g.,
+            ```target_theta = (0.5, 0.5)```.  One of ``target_image`` or ``target_theta`` must be set.
+            Some example values: ``(10.0, 20.0)`` or ``(0.0, 30.0)`` represent T1w;
+            ``(-18.0, -16.0)`` represents T2W; ``(0.0, 0.0)`` represents FLAIR.
+        norm_val: normalization value. Defaults to None.
+        intermediate_out_dir: directory to save intermediate results. Defaults to None.
+        gpu_id: integer number specifies which GPU to run HACA3. Defaults to None.
+        num_batches: During inference, HACA3 takes entire 3D MRI volumes as input.
+            This may cause a considerable amount GPU memory. For reduced GPU memory consumption,
+            source images maybe divided into smaller batches. However, this may slightly
+            increase the inference time. Defaults to None.
+
+    Raises:
+        RuntimeError: _description_
+        RuntimeError: _description_
+
+    Returns:
+        _description_
+    """
+
+    if all(i is None for i in [target_image, target_theta]):
+        raise RuntimeError("Either target_image or target_theta must be set")
+
+    if all(i is not None for i in [target_image, target_theta]):
+        raise RuntimeError("Only one of target_image or target_theta must be set")
+
+    if isinstance(inputs, (str, os.PathLike)) or not isinstance(inputs, Sequence):
+        inputs = (inputs, )
+
+    inputs = [tositk(img) for img in inputs]
+    if target_image is not None: target_image = tositk(target_image)
+
+    # --------------------------------- run HACA3 -------------------------------- #
+    with tempfile.TemporaryDirectory() as tmpdir:
+        for i, img in enumerate(inputs):
+            if tuple(img.GetSize()) != (192, 224, 192):
+                raise RuntimeError(
+                    "all inputs to HACA3 must be in MNI152 space and center-padded to size of ``[192, 224, 192]``. "
+                    f"Got image {i} of size {img.GetSize()}")
+
+            sitk.WriteImage(img, os.path.join(tmpdir, f"input_{i}.nii.gz"))
+
+        if target_image is not None:
+            target_path = os.path.join(tmpdir, "target_image.nii.gz")
+            if tuple(target_image.GetSize()) != (192, 224, 192):
+                raise RuntimeError(
+                    "all inputs to HACA3 must be in MNI152 space and center-padded to size of ``[192, 224, 192]``. "
+                    f"Got target_image of size {target_image.GetSize()}")
+
+            sitk.WriteImage(target_image, target_path)
+        else:
+            target_path = None
+
+        run_HACA3(
+            conda_path=conda_path,
+            env_name=env_name,
+            harmonization_model=harmonization_model,
+            fusion_model=fusion_model,
+            in_path = [os.path.join(tmpdir, f"input_{i}.nii.gz") for i in range(len(inputs))],
+            out_path = os.path.join(tmpdir, "output.nii.gz"),
+            target_image = target_path,
+            target_theta = target_theta,
+            norm_val = norm_val,
+            intermediate_out_dir = intermediate_out_dir,
+            gpu_id = gpu_id,
+            num_batches = num_batches,
+        )
+
+        harmonized = tositk(os.path.join(tmpdir, "output_harmonized_fusion.nii.gz"))
+
+    return harmonized

mrid_python-0.1.7/mrid/preprocessing/spatial.py

@@ -0,0 +1,51 @@
+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`` (numpy axis order, i.e. reversed of sitk size).
+    Retains correct spatial information: origin and direction are preserved and spacing is
+    scaled so the physical field of view stays constant."""
+    img = tositk(img)
+    new_size = list(reversed(new_size))
+
+    old_size = img.GetSize()
+    old_spacing = img.GetSpacing()
+    new_spacing = []
+    for n, o, s in zip(new_size, old_size, old_spacing):
+        if n > 1:
+            new_spacing.append((o - 1) * s / (n - 1))
+        else:
+            new_spacing.append(s)
+
+    reference_image = sitk.Image(new_size, img.GetPixelID())
+    reference_image.SetOrigin(img.GetOrigin())
+    reference_image.SetSpacing(new_spacing)
+    reference_image.SetDirection(img.GetDirection())
+
+    return sitk.Resample(img, reference_image, sitk.Transform(), interpolator, 0.0)
+
+def downsample(image:ImageLike, factor:float, dims: int | Sequence[int] | None, interpolator=sitk.sitkLinear) -> sitk.Image:
+    """factor = 2 for 2x downsampling"""
+    if isinstance(dims, int): dims = (dims, )
+    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_python-0.1.5 → mrid_python-0.1.7}/mrid/study.py

@@ -13,6 +13,7 @@ import SimpleITK as sitk
 
 from . import preprocessing
 from .loading.convert import ImageLike, tonumpy, tositk, totensor
+from .utils.sitk_utils import sitk_apply_numpy
 from .utils.torch_utils import CUDA_IF_AVAILABLE
 
 if TYPE_CHECKING:
@@ -125,6 +126,26 @@ class Study(UserDict[str, sitk.Image | Any]):
 
         return Study(**scans, **seg, **self.get_info())
 
+    def apply_numpy(self, fn:Callable[[np.ndarray], np.ndarray] | None, seg_fn: Callable[[np.ndarray], np.ndarray] | None) -> "Study":
+        """Returns a new ``Study`` with ``fn`` applied to scans and ``seg_fn`` applied to segmentations.
+
+        Args:
+            fn: Function to apply to scan images. Must take and return ``sitk.Image``.
+                If None, identity function is used.
+            seg_fn: Function to apply to segmentation images. Must take and return ``sitk.Image``.
+                If None, identity function is used.
+        """
+        scans = self.get_scans()
+        seg = self.get_segmentations()
+
+        if fn is not None:
+            scans = {k: sitk_apply_numpy(v, fn) for k,v in scans.items()}
+
+        if seg_fn is not None:
+            seg = {k: sitk_apply_numpy(v, seg_fn) for k,v in seg.items()}
+
+        return Study(**scans, **seg, **self.get_info())
+
     def cast(self, dtype) -> "Study":
         """Returns a new study with all scans cast to the specified SimpleITK dtype.
 
@@ -176,6 +197,20 @@ class Study(UserDict[str, sitk.Image | Any]):
         d = preprocessing.cropping.crop_bg_D(self.get_images(), key)
         return Study(**d, **self.get_info())
 
+    def center_crop_or_pad(self, size: Sequence[int]):
+        """Returns a new study with all images cropped or padded from the center to ``size``.
+
+        Args:
+            size: target image shape in pixels.
+        """
+        shapes = {k: tuple(img.GetSize()) for k, img in self.get_images().items()}
+        if len(set(shapes.values())) > 1:
+            raise RuntimeError(f"center_crop_or_pad can only be applied to a Study where all images have the same shape. "
+                               f"Current shapes: {shapes}")
+
+        fn = partial(preprocessing.center_crop_or_pad, size=size)
+        return self.apply(fn=fn, seg_fn=fn)
+
     def skullstrip_hd_bet(
         self,
         key: str,
@@ -213,7 +248,7 @@ class Study(UserDict[str, sitk.Image | Any]):
                 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_hd_bet"`` with brain mask predicted by HD-BET to returned dictionary.
+                if True, adds ``"seg_hd_bet"`` with brain mask predicted by HD-BET to returned study.
                 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 returned study with ``"_hd_bet"`` postfix,
@@ -279,6 +314,85 @@ class Study(UserDict[str, sitk.Image | Any]):
         )
         return Study(**d, **self.get_segmentations(), **self.get_info())
 
+    def harmonize_haca3(
+        self,
+        conda_path: str | os.PathLike,
+        env_name: str,
+        harmonization_model: str | os.PathLike,
+        fusion_model: str | os.PathLike,
+        keys: "str | Sequence[str]",
+        target_image: "ImageLike | None" = None,
+        target_theta: tuple[float,float] | None = None,
+        norm_val: float | None = None,
+        intermediate_out_dir: str | os.PathLike | None = None,
+        gpu_id: int | None = None,
+        num_batches: int | None = None,
+        new_key="harmonized_haca3",
+    ):
+        """Returns a new study, where a harmonized image is added under key ``new_key``
+        (by default ``"harmonized_haca3"``).
+
+        Important: Some preprocessing steps are needed before running HACA3:
+
+        - Inhomogeneity correction
+        - Super-resolution for 2D acquired scans. This step is optional, but recommended for optimal performance. See [SMORE](https://github.com/volcanofly/SMORE-Super-resolution-for-3D-medical-images-MRI) for more details.
+        - Registration to MNI space (1mm isotropic resolution). HACA3 assumes a spatial dimension of 192x224x192.
+
+        For example, you can do the following:
+        ```python
+        study = study.register_SE("t1", mrid.atlas.get_mni152("2006 T1w symmetric"))
+        study = study.n4_bias_field_correction("t1")
+        study = study.center_crop_or_pad([192, 224, 192])
+        study = study.harmonize_haca3(...)
+        ```
+
+        Args:
+            conda_path: path to ``minconda3`` directory.
+            env_name: name of the conda env where HACA3 is installed.
+            harmonization_model: pretrained HACA3 weights. Pretrained model weights on IXI, OASIS and HCP
+                data can be downloaded [here](https://iacl.ece.jhu.edu/~lianrui/haca3/harmonization_public.pt).
+            fusion_model: pretrained fusion model weights. HACA3 uses a 3D convolutional network to
+                combine multi-orientation 2D slices into a single 3D volume. Pretrained fusion model
+                can be downloaded [here](https://iacl.ece.jhu.edu/~lianrui/haca3/fusion.pt).
+            keys: key of the input source image. Multiple key may be provided if there are
+                multiple source images (different modalities). Note that modalities must be in
+                MNI space (1mm isotropic resolution). HACA3 assumes a spatial dimension of 192x224x192.
+            target_image: target image. HACA3 will match the contrast of
+                source images to this target image. One of ``target_image`` or ``target_theta`` must be set.
+            target_theta: In HACA3, ``theta`` is a two-dimensional representation of image contrast.
+                Target image contrast can be directly specified by providing a ``theta`` value, e.g.,
+                ``target_theta = (0.5, 0.5)``.  One of ``target_image`` or ``target_theta`` must be set.
+                Some example values: ``(10.0, 20.0)`` or ``(0.0, 30.0)`` represent T1w;
+                ``(-18.0, -16.0)`` represents T2W; ``(0.0, 0.0)`` represents FLAIR.
+            norm_val: normalization value. Defaults to None.
+            intermediate_out_dir: directory to save intermediate results. Defaults to None.
+            gpu_id: integer number specifies which GPU to run HACA3. Defaults to None.
+            num_batches: During inference, HACA3 takes entire 3D MRI volumes as input.
+                This may cause a considerable amount GPU memory. For reduced GPU memory consumption,
+                source images maybe divided into smaller batches. However, this may slightly
+                increase the inference time. Defaults to None.
+            new_key: name of the harmonized image in retuned study. Defaults to "harmonized_haca3".
+        """
+        if isinstance(keys, str): keys = (keys, )
+        images = [(k, self[k]) for k in keys]
+
+        harmonized = preprocessing.haca3.harmonize(
+            conda_path=conda_path,
+            env_name=env_name,
+            harmonization_model=harmonization_model,
+            fusion_model=fusion_model,
+            inputs=[img for k,img in images],
+            target_image=target_image,
+            target_theta=target_theta,
+            norm_val=norm_val,
+            intermediate_out_dir=intermediate_out_dir,
+            gpu_id=gpu_id,
+            num_batches=num_batches,
+        )
+
+        return self.add(new_key, harmonized)
+
+
     def resize(self, size: Sequence[int], interpolator=sitk.sitkLinear):
         """Returns a new study with all images resized to to ``size``.
 
@@ -428,7 +542,7 @@ class Study(UserDict[str, sitk.Image | Any]):
                 if specified, processed mask is added to returned study with specified postfix rather than
                 replacing current ``key``.
         """
-        arr = self.numpy(key)
+        arr = self.to_numpy(key)
 
         if independent_channels:
             from monai.transforms import remove_small_objects
@@ -473,14 +587,14 @@ class Study(UserDict[str, sitk.Image | Any]):
         from monai.transforms import KeepLargestConnectedComponent
         tfm = KeepLargestConnectedComponent(applied_labels=applied_labels, is_onehot=False, independent=independent,
                                             connectivity=connectivity, num_components=num_components)
-        arr = tfm(self.numpy(key))
+        arr = tfm(self.to_numpy(key))
         return self.add(f'{key}{postfix}', arr, reference_key=key)
 
-    def numpy(self, key: str):
+    def to_numpy(self, key: str):
         """returns ``study[key]`` converted to a numpy array."""
         return tonumpy(self[key])
 
-    def tensor(self, key: str):
+    def to_tensor(self, key: str):
         """returns ``study[key]`` converted to a tensor."""
         return totensor(self[key])
 
@@ -531,18 +645,18 @@ class Study(UserDict[str, sitk.Image | Any]):
         stacked = torch.stack([torch.from_numpy(sitk.GetArrayFromImage(v)) for _,v in items])
         return stacked.to(device=device, dtype=dtype, memory_format=torch.contiguous_format)
 
-    def numpy_dict(self) -> dict[str, np.ndarray | Any]:
+    def to_numpy_dict(self) -> dict[str, np.ndarray | Any]:
         """Returns a dictionary with all images converted to numpy arrays, info is included as is."""
         return {k: (sitk.GetArrayFromImage(v) if isinstance(v, sitk.Image) else v) for k, v in self.items()}
 
-    def tensor_dict(self) -> "dict[str, torch.Tensor | Any]":
+    def to_tensor_dict(self) -> "dict[str, torch.Tensor | Any]":
         """Returns a dictionary with all images converted to tensors, info is included as is."""
         import torch
-        return {k: (torch.from_numpy(v) if isinstance(v, np.ndarray) else v) for k,v in self.numpy_dict()}
+        return {k: (torch.from_numpy(v) if isinstance(v, np.ndarray) else v) for k,v in self.to_numpy_dict()}
 
     def plot(self):
         from .utils.plotting import plot_study
-        return plot_study(self.get_images().numpy_dict())
+        return plot_study(self.get_images().to_numpy_dict())
 
     def save(
         self,

mrid_python-0.1.7/mrid/training/cropping.py

@@ -0,0 +1,52 @@
+import random
+from collections.abc import Sequence, Callable
+from typing import Any, Literal, TypeVar
+
+ARR = TypeVar("ARR", bound=Any)
+
+def crop(
+    arr: ARR,
+    reduction: Sequence[int],
+    where: Literal["start", "end", "center", "random"] = "center",
+) -> ARR:
+    """Crop ``arr`` such that ``output.shape[i] = input.shape[i] - reduction[i]``"""
+
+    shape = arr.shape[-len(reduction):]
+    slices = []
+
+    for r, sh in zip(reduction, shape):
+        if r == 0:
+            slices.append(slice(None))
+            continue
+
+        if r < 0: raise ValueError(f"Reduction {r} cannot be negative")
+        if r > sh: raise ValueError(f"Reduction {r} exceeds dimension size {sh}")
+
+        if where == 'start': start, end = 0, sh - r
+        elif where == 'end': start, end = r, sh
+        elif where == 'center':
+            start = r // 2
+            end = start + (sh - r)
+        elif where == 'random':
+            start = random.randint(0, r)
+            end = start + (sh - r)
+        else:
+            raise ValueError(f"Invalid where: {where}")
+
+        slices.append(slice(start, end))
+
+    # apply with broadcasting
+    return arr[(..., *slices)]
+
+def crop_to_shape(
+    input: ARR,
+    shape: Sequence[int],
+    where: Literal["start", "end", "center", "random"] = "center",
+) -> ARR:
+    """Crop ``input`` to ``shape``."""
+
+    # broadcast
+    if len(shape) < input.ndim:
+        shape = list(input.shape[:input.ndim - len(shape)]) + list(shape)
+
+    return crop(input, [i - j for i, j in zip(input.shape, shape)], where=where)

mrid_python-0.1.7/mrid/training/padding.py

@@ -0,0 +1,97 @@
+import random
+from collections.abc import Sequence, Callable
+from typing import Any, Literal, TypeVar
+
+import numpy as np
+import torch
+from . import cropping
+from ..utils.python_utils import reduce_dim
+
+ARR = TypeVar("ARR", bound=Any)
+
+@torch.no_grad
+def pad(
+    input: ARR,
+    padding: Sequence[int],
+    mode: str = "constant",
+    value = None,
+    where: Literal["center", "start", "end"] = "center",
+    crop: bool = False,
+) -> ARR:
+    """
+    TODO REFACTOR
+
+    ``output.shape[i] = input.shape[i] + padding[i]``.
+
+    Args:
+        input (torch.Tensor): input to pad.
+        padding (str): how much padding to add per each dimension of ``input``.
+        mode (str, optional):
+            padding mode (https://pytorch.org/docs/stable/generated/torch.nn.functional.pad.html).
+            Defaults to 'constant'.
+        value (_type_, optional): padding constant value. Defaults to None.
+        where (str, optional): where to pad.
+            if ``center``, will pad start and end of each dimension evenly,
+            if ``start``, will pad at the start of each dimension,
+            if ``end``, will pad at the end. Defaults to 'center'.
+        crop (bool, optional): allow cropping if padding is negative. Defaults to False.
+
+    Returns:
+        torch.Tensor: Padded `input`.
+    """
+    pad_values = [i if i > 0 else 0 for i in padding]
+
+    if sum(pad_values) > 0:
+
+        # create padding sequence for torch.nn.functional.pad
+        if where == 'center':
+            dims_padding = [(int(i / 2), int(i / 2)) if i % 2 == 0 else (int(i / 2), int(i / 2) + 1) for i in padding]
+        elif where == 'start':
+            dims_padding = [(i, 0) for i in padding]
+        elif where == 'end':
+            dims_padding = [(0, i) for i in padding]
+        else: raise ValueError(f'Invalid where: {where}')
+
+        # broadcasting (e.g. if padding 3×128×128 by [4, 4], it will pad by [0, 4, 4])
+        if len(dims_padding) < input.ndim:
+            dims_padding = [(0, 0)] * (input.ndim - len(dims_padding)) + dims_padding
+
+        if mode == 'zeros': mode = 'constant'; value = 0
+        elif mode == 'min': mode = 'constant'; value = float(input.min())
+        elif mode == 'max': mode = 'constant'; value = float(input.max())
+        elif mode == 'mean': mode = 'constant'; value = float(input.mean())
+
+        if isinstance(input, np.ndarray):
+            if mode == 'constant': kwargs = {"constant_values": value}
+            else: kwargs = {}
+            input = np.pad(input, pad_width = dims_padding, mode = mode, **kwargs) # type:ignore
+        else:
+            input = torch.nn.functional.pad(input, reduce_dim(reversed(dims_padding)), mode=mode, value=value) # type:ignore
+
+    if crop:
+        crop_values = [-i if i < 0 else 0 for i in padding]
+        if sum(crop_values) > 0:
+            input =  cropping.crop(input, crop_values, where=where)
+
+    return input
+
+def pad_to_shape(
+    input:ARR,
+    shape:Sequence[int],
+    mode:str = "constant",
+    value=None,
+    where:Literal["center", "start", "end"] = "center",
+    crop = False,
+) -> ARR:
+    # broadcasting
+    if len(shape) < input.ndim:
+        shape = list(input.shape[:input.ndim - len(shape)]) + list(shape)
+
+    return pad(
+        input=input,
+        padding=[shape[i] - input.shape[i] for i in range(input.ndim)],
+        mode=mode,
+        value=value,
+        where=where,
+        crop = crop,
+    )

{mrid_python-0.1.5 → mrid_python-0.1.7}/mrid/training/slicer.py

@@ -43,7 +43,7 @@ class SliceSampler:
             raise RuntimeError(f"Segmentation must have a shape of of (D, H, W), got {segmentation.shape}")
         if segmentation.is_floating_point():
             raise RuntimeError(f"Segmentation must have integer data type, got {segmentation.dtype}")
-        if segmentation.shape[1:] != data.shape:
+        if data.shape[1:] != segmentation.shape:
             raise RuntimeError(f"Shapes of scans and segmentation do not match: {data.shape = }, {segmentation.shape = }")
         if segmentation.min() < 0:
             raise RuntimeError(f"Segmentation background must have a value of 0, got {segmentation.min() = }")
@@ -193,6 +193,11 @@ class SliceSampler:
 
         return get_sample
 
+    def get_all_dim_slices(self, dim: Literal[0,1,2], around: int, flatten: bool = True):
+        return [self.get_slice(
+            dim=dim, coord=coord, around=around, flatten=flatten
+        ) for coord in range(around, self.data.shape[dim+1] - around)]
+
 
 class SliceDataset(torch.utils.data.Dataset):
     """A dataset of SliceSamplers.

{mrid_python-0.1.5 → mrid_python-0.1.7}/mrid/training/transforms.py

@@ -2,58 +2,9 @@ import random
 from collections.abc import Sequence, Callable
 from typing import Any, Literal, TypeVar
 
+import numpy as np
 import torch
 
-ARR = TypeVar("ARR", bound=Any)
-
-def crop(
-    arr: ARR,
-    reduction: Sequence[int],
-    where: Literal["start", "end", "center", "random"] = "center",
-) -> ARR:
-    """Crop ``arr`` such that ``output.shape[i] = input.shape[i] - reduction[i]``"""
-
-    shape = arr.shape[-len(reduction):]
-    slices = []
-
-    for r, sh in zip(reduction, shape):
-        if r == 0:
-            slices.append(slice(None))
-            continue
-
-        if r < 0: raise ValueError("Reduction cannot be negative")
-        if r > sh: raise ValueError(f"Reduction {r} exceeds dimension size {sh}")
-
-        if where == 'start': start, end = 0, sh - r
-        elif where == 'end': start, end = r, sh
-        elif where == 'center':
-            start = r // 2
-            end = start + (sh - r)
-        elif where == 'random':
-            start = random.randint(0, r)
-            end = start + (sh - r)
-        else:
-            raise ValueError(f"Invalid where: {where}")
-
-        slices.append(slice(start, end))
-
-    # apply with broadcasting
-    return arr[(..., *slices)]
-
-def crop_to_shape(
-    input: ARR,
-    shape: Sequence[int],
-    where: Literal["start", "end", "center", "random"] = "center",
-) -> ARR:
-    """Crop ``input`` to ``shape``."""
-
-    # broadcast
-    if len(shape) < input.ndim:
-        shape = list(input.shape[:input.ndim - len(shape)]) + list(shape)
-
-    return crop(input, [i - j for i, j in zip(input.shape, shape)], where=where)
-
-
 def shuffle_channels(x:torch.Tensor):
     """Shuffle first axis in a ``(C, *)`` tensor"""
     return x[torch.randperm(x.shape[0])]

{mrid_python-0.1.5 → mrid_python-0.1.7}/mrid/utils/dcm2niix.py

@@ -47,10 +47,12 @@ def run_dcm2niix(
 
         # create temporary folders
         tmp_input_dir = os.path.join(tmpdir, "mrid_dcm2niix_input")
+        if os.path.exists(tmp_input_dir): shutil.rmtree(tmp_input_dir)
+        shutil.copytree(inpath, tmp_input_dir)
+
         tmp_output_dir = os.path.join(tmpdir, "mrid_dcm2niix_output")
         if os.path.exists(tmp_output_dir): shutil.rmtree(tmp_output_dir)
         os.mkdir(tmp_output_dir)
-        shutil.copytree(inpath, tmp_input_dir)
 
         # create output dir if it doesn't exist
         if outfolder != '':
@@ -72,7 +74,7 @@ def run_dcm2niix(
         out_files = [i for i in os.listdir(tmp_output_dir) if i.lower().strip().endswith('.nii.gz')]
 
         # move them to output folder
-        shutil.copytree(tmp_output_dir, outfolder)
+        shutil.copytree(tmp_output_dir, outfolder, dirs_exist_ok=True)
 
         if len(out_files) > 1:
             warnings.warn(f"More than one NIfTI file was created in {outfolder}, path to the first one will be returned. Something may be wrong.")

{mrid_python-0.1.5 → mrid_python-0.1.7}/mrid/utils/plotting.py

@@ -1,12 +1,16 @@
 from collections.abc import Mapping
 
 import numpy as np
+
 from ..loading import ImageLike, tonumpy
 
-def plot_study(data: Mapping[str, ImageLike]):
+def plot_study(data: "ImageLike | Mapping[str, ImageLike]"):
     import matplotlib.gridspec as gridspec
     import matplotlib.pyplot as plt
 
+    if not isinstance(data, Mapping):
+        data = {"image": data}
+
     data = {k: tonumpy(v) for k,v in data.items()}
     n_vals = len(data)
 
@@ -78,29 +82,28 @@ def plot_study(data: Mapping[str, ImageLike]):
         fig.text(box.x0 + box.width/2, box.y1 + 0.04, modality_name,
                  ha='center', va='bottom', fontsize=14, fontweight='bold')
 
-    # plt.show()
     return fig
 
-if __name__ == "__main__":
-
-    def dummy_data(shape, type='sphere'):
-        x, y, z = np.indices(shape)
-        cx, cy, cz = shape[0]//2, shape[1]//2, shape[2]//2
-        mask = np.zeros(shape)
-        if type == 'sphere':
-            r = min(shape)//3
-            mask[(x-cx)**2 + (y-cy)**2 + (z-cz)**2 < r**2] = 1
-        elif type == 'cube':
-            r = min(shape)//4
-            mask[cx-r:cx+r, cy-r:cy+r, cz-r:cz+r] = 1
-        elif type == 'noise':
-            mask = np.random.rand(*shape)
-
-        return mask
-
-    plot_study({
-        "Sphere": dummy_data((60, 60, 60), 'sphere'),
-        "Cube": dummy_data((60, 60, 60), 'cube'),
-        "Noise": dummy_data((60, 60, 60), 'noise'),
-        "Another sphere": dummy_data((60, 60, 60), 'sphere'),
-    })
+# if __name__ == "__main__":
+
+#     def dummy_data(shape, type='sphere'):
+#         x, y, z = np.indices(shape)
+#         cx, cy, cz = shape[0]//2, shape[1]//2, shape[2]//2
+#         mask = np.zeros(shape)
+#         if type == 'sphere':
+#             r = min(shape)//3
+#             mask[(x-cx)**2 + (y-cy)**2 + (z-cz)**2 < r**2] = 1
+#         elif type == 'cube':
+#             r = min(shape)//4
+#             mask[cx-r:cx+r, cy-r:cy+r, cz-r:cz+r] = 1
+#         elif type == 'noise':
+#             mask = np.random.rand(*shape)
+
+#         return mask
+
+#     plot_study({
+#         "Sphere": dummy_data((60, 60, 60), 'sphere'),
+#         "Cube": dummy_data((60, 60, 60), 'cube'),
+#         "Noise": dummy_data((60, 60, 60), 'noise'),
+#         "Another sphere": dummy_data((60, 60, 60), 'sphere'),
+#     })

{mrid_python-0.1.5 → mrid_python-0.1.7}/mrid/utils/python_utils.py

@@ -1,7 +1,7 @@
 import importlib.util
-from collections.abc import Mapping, Sequence
-from typing import Any
-
+from collections.abc import Mapping, Sequence, Iterable
+from typing import Any, TypeVar
+import functools, operator
 
 # lazy loader from https://stackoverflow.com/a/78312674/15673832
 class LazyLoader:
@@ -30,19 +30,7 @@ class LazyLoader:
         return getattr (self._mod, attr)
 
 
-# this allows transforms to support any kind of container
-# class _Packer[T: Sequence | Mapping]:
-#     def __init__(self, type: type[T], keys: list | None = None):
-#         self.type: Any = type
-#         self.keys = keys
-
-#     def pack(self, unpacked: Sequence) -> T:
-#         if self.keys is not None: return self.type(dict(zip(self.keys, unpacked)))
-#         return self.type(unpacked)
-
-# def unpack_struct[T: Sequence | Mapping](struct: T) -> tuple[Any, _Packer[T]]:
-#     if isinstance(struct, Sequence):
-#         return list(struct), _Packer(type(struct))
-#     if isinstance(struct, Mapping):
-#         return list(struct.values()), _Packer(type(struct), list(struct.values()))
-#     raise TypeError(f"Transformation functions accept lists and dictionaries, but received {type(struct)}")
+T = TypeVar("T")
+def reduce_dim(x:Iterable[Iterable[T]]) -> list[T]:
+    """Reduces one level of nesting. Takes an iterable of iterables of X, and returns an iterable of X."""
+    return functools.reduce(operator.iconcat, x, [])

mrid_python-0.1.7/mrid/utils/sitk_utils.py

@@ -0,0 +1,18 @@
+from collections.abc import Callable
+
+import numpy as np
+import SimpleITK as sitk
+
+from ..loading import ImageLike, tositk
+
+
+def sitk_apply_numpy(image: ImageLike, func: Callable[[np.ndarray], np.ndarray]) -> sitk.Image:
+    image = tositk(image)
+    array = sitk.GetArrayFromImage(image)
+    shape = array.shape
+    array = func(array)
+    if shape != array.shape:
+        raise RuntimeError(f"Function {func} changed array shape from {shape} to {array.shape}.")
+    res = sitk.GetImageFromArray(array)
+    res.CopyInformation(image)
+    return res

{mrid_python-0.1.5 → mrid_python-0.1.7}/mrid_python.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mrid-python
-Version: 0.1.5
+Version: 0.1.7
 Summary: Tools for working with 3D medical images and segmentations - registration, brain skull-stripping, etc.
 Author-email: Ivan Nikishev <nkshv2@gmail.com>
 Project-URL: Homepage, https://github.com/inikishev/mrid

{mrid_python-0.1.5 → mrid_python-0.1.7}/mrid_python.egg-info/SOURCES.txt

@@ -5,18 +5,22 @@ mrid/study.py
 mrid/atlas/__init__.py
 mrid/atlas/MNI152/__init__.py
 mrid/atlas/SRI24/__init__.py
+mrid/inference/__init__.py
 mrid/loading/__init__.py
 mrid/loading/convert.py
 mrid/preprocessing/CTseg.py
 mrid/preprocessing/__init__.py
 mrid/preprocessing/bias_field_correction.py
 mrid/preprocessing/cropping.py
+mrid/preprocessing/haca3.py
 mrid/preprocessing/hd_bet.py
 mrid/preprocessing/mask.py
 mrid/preprocessing/simple_elastix.py
 mrid/preprocessing/spatial.py
 mrid/preprocessing/synthstrip.py
 mrid/training/__init__.py
+mrid/training/cropping.py
+mrid/training/padding.py
 mrid/training/slicer.py
 mrid/training/transforms.py
 mrid/utils/__init__.py
@@ -24,6 +28,7 @@ mrid/utils/dcm2niix.py
 mrid/utils/dicom_uid_fixer.py
 mrid/utils/plotting.py
 mrid/utils/python_utils.py
+mrid/utils/sitk_utils.py
 mrid/utils/stl_utils.py
 mrid/utils/torch_utils.py
 mrid_python.egg-info/PKG-INFO

{mrid_python-0.1.5 → mrid_python-0.1.7}/pyproject.toml

@@ -13,7 +13,7 @@ build-backend = "setuptools.build_meta"
 name = "mrid-python"
 description = "Tools for working with 3D medical images and segmentations - registration, brain skull-stripping, etc."
 
-version = "0.1.5"
+version = "0.1.7"
 dependencies = [
   "numpy",
   "SimpleITK",

{mrid_python-0.1.5 → mrid_python-0.1.7}/tests/test_preprocessing.py

@@ -10,8 +10,8 @@ def test_resize():
     study = Study(t1=np.random.rand(10, 20, 30).astype(np.float32))
     resized = study.resize([5, 10, 15])
     assert 't1' in resized
-    assert study.numpy("t1").shape == (10, 20, 30)
-    assert resized.numpy("t1").shape == (5, 10, 15)
+    assert study.to_numpy("t1").shape == (10, 20, 30)
+    assert resized.to_numpy("t1").shape == (5, 10, 15)
 
 
 def test_downsample():
@@ -19,8 +19,8 @@ def test_downsample():
     study = Study(t1=np.random.rand(10, 20, 30).astype(np.float32))
     downsampled = study.downsample(factor=2.0)  # 2x downsampling
     assert 't1' in downsampled
-    assert study.numpy("t1").shape == (10, 20, 30)
-    assert downsampled.numpy("t1").shape == (5, 10, 15)
+    assert study.to_numpy("t1").shape == (10, 20, 30)
+    assert downsampled.to_numpy("t1").shape == (5, 10, 15)
 
 
 def test_bias_field_correction():
@@ -40,4 +40,4 @@ def test_crop_bg():
     assert 't1' in cropped
     assert 't2' in cropped
 
-    assert study.numpy("t1").shape == study.numpy("t2").shape
+    assert study.to_numpy("t1").shape == study.to_numpy("t2").shape

{mrid_python-0.1.5 → mrid_python-0.1.7}/tests/test_study.py

@@ -90,7 +90,7 @@ def test_numpy_method():
     data = np.random.rand(10, 20, 30).astype(np.float32)
     study = Study(t1=data)
 
-    numpy_array = study.numpy('t1')
+    numpy_array = study.to_numpy('t1')
     assert isinstance(numpy_array, np.ndarray)
     assert numpy_array.shape == (10, 20, 30)
 

mrid_python-0.1.5/mrid/preprocessing/spatial.py

@@ -1,82 +0,0 @@
-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: int | Sequence[int] | None, interpolator=sitk.sitkLinear) -> sitk.Image:
-    """factor = 2 for 2x downsampling"""
-    if isinstance(dims, int): dims = (dims, )
-    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)