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

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/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/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/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/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/preprocessing/hd_bet.py

@@ -235,7 +235,7 @@ def skullstrip_D(
 
     # optionally add with skullstripped postfix
     if keep_original:
-        skullstripped = {f"{k}_hd_bet": v for k,v in skullstripped}
+        skullstripped = {(f"{k}_hd_bet" if k != "seg_hd_bet" else k): v for k,v in skullstripped.items()}
         skullstripped.update(images.copy())
 
     return skullstripped

mrid/preprocessing/mask.py

@@ -1,6 +1,5 @@
 import SimpleITK as sitk
 import numpy as np
-import SimpleITK as sitk
 
 from ..loading.convert import ImageLike, tositk, tonumpy
 
@@ -15,12 +14,12 @@ def expand_binary_mask(binary_mask: ImageLike, expand: int) -> sitk.Image:
             Negative values dilate the mask by this many pixels.
     """
     binary_mask = tositk(binary_mask)
-    if expand > 0:
+    if expand < 0:
         inverted_mask = 1 - binary_mask
-        return 1 - sitk.BinaryDilate(inverted_mask, (expand, expand, expand))
+        return 1 - sitk.BinaryDilate(inverted_mask, (-expand, -expand, -expand))
 
-    if expand < 0:
-        return sitk.BinaryDilate(binary_mask, (-expand, -expand, -expand))
+    if expand > 0:
+        return sitk.BinaryDilate(binary_mask, (expand, expand, expand))
 
     return binary_mask
 

mrid/preprocessing/spatial.py

@@ -20,61 +20,31 @@ def resample_to(input: ImageLike, to: ImageLike, interpolation=sitk.sitkNearestN
 
 
 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"""
+    """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))
 
-    # img = sitk.ReadImage(img)
-    dimension = img.GetDimension()
+    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)
 
-    # 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_image = sitk.Image(new_size, img.GetPixelID())
+    reference_image.SetOrigin(img.GetOrigin())
+    reference_image.SetSpacing(new_spacing)
+    reference_image.SetDirection(img.GetDirection())
 
-    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)]
+    return sitk.Resample(img, reference_image, sitk.Transform(), interpolator, 0.0)
 
-    # 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:
+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)]

mrid/preprocessing/synthstrip.py

@@ -70,6 +70,7 @@ def run_synthstrip(
     fill: int | None = None,
     no_csf: bool | None = None,
     model: str | os.PathLike | None = None,
+    verbose: bool = True,
 ):
     """Runs ``synthstrip`` command-line routine via ``subprocess.run``.
 
@@ -108,7 +109,10 @@ def run_synthstrip(
     if model is not None: command.extend(["--model", os.path.normpath(model)])
 
     # run
-    subprocess.run(command, check=True)
+    if verbose:
+        subprocess.run(command, check=True)
+    else:
+        subprocess.run(command, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT)
 
 def predict_brain_mask(
     synthstrip_script_path: str | os.PathLike,
@@ -117,6 +121,7 @@ def predict_brain_mask(
     border: int | None = None,
     threads: int | None = None,
     model: str | os.PathLike | None = None,
+    verbose: bool = True,
 ):
     """Returns brain mask of ``input`` predicted by ``synthstrip``.
 
@@ -141,6 +146,7 @@ def predict_brain_mask(
             border=border,
             threads=threads,
             model=model,
+            verbose=verbose,
         )
 
         brain_mask = tositk(os.path.join(tmpdir, "synthstrip_mask.nii.gz"))
@@ -155,6 +161,7 @@ def skullstrip(
     threads: int | None = None,
     model: str | os.PathLike | None = None,
     expand: int = 0,
+    verbose: bool = True,
 ):
     """Skullstrips ``input`` using synthstrip.
 
@@ -181,6 +188,7 @@ def skullstrip(
         border=border,
         threads=threads,
         model=model,
+        verbose=verbose
     )
     if expand != 0:
         mask = expand_binary_mask(mask, expand=expand)
@@ -202,6 +210,8 @@ def skullstrip_D(
     include_mask: bool = False,
     keep_original: bool = False,
 
+    verbose: bool = True,
+
 ) -> dict[str, sitk.Image]:
     """Predicts brain mask of ``images[key]`` using synthstrip, then uses this mask to skull strip all values in ``images``.
 
@@ -233,6 +243,7 @@ def skullstrip_D(
         border=border,
         threads=threads,
         model=model,
+        verbose=verbose,
     )
     skullstripped = {}
 
@@ -252,7 +263,7 @@ def skullstrip_D(
 
     # optionally add with skullstripped postfix
     if keep_original:
-        skullstripped = {f"{k}_synthstrip": v for k,v in skullstripped}
+        skullstripped = {(f"{k}_synthstrip" if k != "seg_synthstrip" else k): v for k,v in skullstripped.items()}
         skullstripped.update(images.copy())
 
     return skullstripped