junifer 0.0.7.dev332__py3-none-any.whl → 0.0.8__py3-none-any.whl

junifer/preprocess/warping/_ants_warper.py

@@ -12,9 +12,11 @@ import nibabel as nib
 import numpy as np
 
 from ...data import get_template, get_xfm
-from ...pipeline import WorkDirManager
-from ...typing import Dependencies, ExternalDependencies
-from ...utils import logger, raise_error, run_ext_cmd
+from ...pipeline import ExtDep, WorkDirManager
+from ...typing import Dependencies, ExternalDependencies, SpaceLike
+from ...utils import raise_error, run_ext_cmd
+from ..base import logger
+from ._utils import warp_in_chunks
 
 
 __all__ = ["ANTsWarper"]
@@ -30,7 +32,7 @@ class ANTsWarper:
 
     _EXT_DEPENDENCIES: ClassVar[ExternalDependencies] = [
         {
-            "name": "ants",
+            "name": ExtDep.ANTs,
             "commands": ["ResampleImage", "antsApplyTransforms"],
         },
     ]
@@ -41,7 +43,7 @@ class ANTsWarper:
         self,
         input: dict[str, Any],
         extra_input: dict[str, Any],
-        reference: str,
+        reference: SpaceLike,
     ) -> dict[str, Any]:
         """Preprocess using ANTs.
 
@@ -52,7 +54,7 @@ class ANTsWarper:
         extra_input : dict
             The other fields in the Junifer Data object. Should have ``T1w``
             and ``Warp`` data types.
-        reference : str
+        reference : str or Enum
             The data type or template space to use as reference for warping.
 
         Returns
@@ -92,7 +94,7 @@ class ANTsWarper:
 
             # Create a tempfile for resampled reference output
             resample_image_out_path = (
-                element_tempdir / "resampled_reference.nii.gz"
+                element_tempdir / "resampled_reference.nii"
             )
             # Set ResampleImage command
             resample_image_cmd = [
@@ -108,21 +110,27 @@ class ANTsWarper:
             run_ext_cmd(name="ResampleImage", cmd=resample_image_cmd)
 
             # Create a tempfile for warped output
-            apply_transforms_out_path = element_tempdir / "warped_data.nii.gz"
-            # Set antsApplyTransforms command
-            apply_transforms_cmd = [
-                "antsApplyTransforms",
-                "-d 3",
-                "-e 3",
-                "-n LanczosWindowedSinc",
-                f"-i {input['path'].resolve()}",
-                # use resampled reference
-                f"-r {resample_image_out_path.resolve()}",
-                f"-t {warp_file_path.resolve()}",
-                f"-o {apply_transforms_out_path.resolve()}",
-            ]
-            # Call antsApplyTransforms
-            run_ext_cmd(name="antsApplyTransforms", cmd=apply_transforms_cmd)
+            apply_transforms_out_path = element_tempdir / "warped_data.nii"
+            # Call antsApplyTransforms in single precision, in chunks of
+            # volumes
+            warp_in_chunks(
+                name="antsApplyTransforms",
+                input_path=input["path"].resolve(),
+                output_path=apply_transforms_out_path.resolve(),
+                cmd=lambda in_path, out_path: [
+                    "antsApplyTransforms",
+                    "-d 3",
+                    "-e 3",
+                    "-n LanczosWindowedSinc",
+                    "--float 1",
+                    f"-i {in_path}",
+                    # use resampled reference
+                    f"-r {resample_image_out_path.resolve()}",
+                    f"-t {warp_file_path.resolve()}",
+                    f"-o {out_path}",
+                ],
+                tempdir=element_tempdir,
+            )
 
             logger.debug("Updating warped data")
             input.update(
@@ -144,7 +152,7 @@ class ANTsWarper:
             if input.get("mask") is not None:
                 # Create a tempfile for warped mask output
                 apply_transforms_mask_out_path = (
-                    element_tempdir / "warped_mask.nii.gz"
+                    element_tempdir / "warped_mask.nii"
                 )
                 # Set antsApplyTransforms command
                 apply_transforms_mask_cmd = [
@@ -152,6 +160,7 @@ class ANTsWarper:
                     "-d 3",
                     "-e 3",
                     "-n 'GenericLabel[NearestNeighbor]'",
+                    "--float 1",
                     f"-i {input['mask']['path'].resolve()}",
                     # use resampled reference
                     f"-r {input['reference']['path'].resolve()}",
@@ -211,7 +220,7 @@ class ANTsWarper:
                     # Create a tempfile for resampled reference output
                     ref_path = (
                         element_tempdir
-                        / f"resampled_reference-{reference}.nii.gz"
+                        / f"resampled_reference-{reference}.nii"
                     )
                     # Set ResampleImage command
                     resample_image_cmd = [
@@ -242,27 +251,33 @@ class ANTsWarper:
                     extra_input=None,
                 )
                 # Save template
-                ref_path = element_tempdir / f"{reference}_T1w.nii.gz"
+                ref_path = element_tempdir / f"{reference}_T1w.nii"
                 nib.save(template_space_img, ref_path)
 
             # Create a tempfile for warped output
             warped_output_path = element_tempdir / (
-                f"warped_data_from_{input_space}_to_{reference}.nii.gz"
+                f"warped_data_from_{input_space}_to_{reference}.nii"
             )
 
-            # Set antsApplyTransforms command
-            apply_transforms_cmd = [
-                "antsApplyTransforms",
-                "-d 3",
-                "-e 3",
-                "-n LanczosWindowedSinc",
-                f"-i {input['path'].resolve()}",
-                f"-r {ref_path.resolve()}",
-                f"-t {xfm_file_path.resolve()}",
-                f"-o {warped_output_path.resolve()}",
-            ]
-            # Call antsApplyTransforms
-            run_ext_cmd(name="antsApplyTransforms", cmd=apply_transforms_cmd)
+            # Call antsApplyTransforms in single precision, in chunks of
+            # volumes
+            warp_in_chunks(
+                name="antsApplyTransforms",
+                input_path=input["path"].resolve(),
+                output_path=warped_output_path.resolve(),
+                cmd=lambda in_path, out_path: [
+                    "antsApplyTransforms",
+                    "-d 3",
+                    "-e 3",
+                    "-n LanczosWindowedSinc",
+                    "--float 1",
+                    f"-i {in_path}",
+                    f"-r {ref_path.resolve()}",
+                    f"-t {xfm_file_path.resolve()}",
+                    f"-o {out_path}",
+                ],
+                tempdir=element_tempdir,
+            )
 
             logger.debug("Updating warped data")
             input.update(
@@ -286,7 +301,7 @@ class ANTsWarper:
                 logger.debug("Warping associated mask")
                 # Create a tempfile for warped mask output
                 apply_transforms_mask_out_path = element_tempdir / (
-                    f"warped_mask_from_{input_space}_to_{reference}.nii.gz"
+                    f"warped_mask_from_{input_space}_to_{reference}.nii"
                 )
                 # Set antsApplyTransforms command
                 apply_transforms_mask_cmd = [
@@ -294,6 +309,7 @@ class ANTsWarper:
                     "-d 3",
                     "-e 3",
                     "-n 'GenericLabel[NearestNeighbor]'",
+                    "--float 1",
                     f"-i {input['mask']['path'].resolve()}",
                     # use resampled reference or original
                     f"-r {ref_path.resolve()}",

junifer/preprocess/warping/_fsl_warper.py

@@ -11,9 +11,11 @@ from typing import (
 import nibabel as nib
 import numpy as np
 
-from ...pipeline import WorkDirManager
+from ...pipeline import ExtDep, WorkDirManager
 from ...typing import Dependencies, ExternalDependencies
-from ...utils import logger, raise_error, run_ext_cmd
+from ...utils import raise_error, run_ext_cmd
+from ..base import logger
+from ._utils import warp_in_chunks
 
 
 __all__ = ["FSLWarper"]
@@ -29,7 +31,7 @@ class FSLWarper:
 
     _EXT_DEPENDENCIES: ClassVar[ExternalDependencies] = [
         {
-            "name": "fsl",
+            "name": ExtDep.FSL,
             "commands": ["flirt", "applywarp"],
         },
     ]
@@ -92,9 +94,10 @@ class FSLWarper:
                 )
 
             # Create a tempfile for resampled reference output
-            flirt_out_path = element_tempdir / "resampled_reference.nii.gz"
+            flirt_out_path = element_tempdir / "resampled_reference.nii"
             # Set flirt command
             flirt_cmd = [
+                "FSLOUTPUTTYPE=NIFTI",
                 "flirt",
                 "-interp spline",
                 f"-in {extra_input['T1w']['path'].resolve()}",
@@ -106,19 +109,24 @@ class FSLWarper:
             run_ext_cmd(name="flirt", cmd=flirt_cmd)
 
             # Create a tempfile for warped output
-            applywarp_out_path = element_tempdir / "warped_data.nii.gz"
-            # Set applywarp command
-            applywarp_cmd = [
-                "applywarp",
-                "--interp=spline",
-                f"-i {input['path'].resolve()}",
-                # use resampled reference
-                f"-r {flirt_out_path.resolve()}",
-                f"-w {warp_file_path.resolve()}",
-                f"-o {applywarp_out_path.resolve()}",
-            ]
-            # Call applywarp
-            run_ext_cmd(name="applywarp", cmd=applywarp_cmd)
+            applywarp_out_path = element_tempdir / "warped_data.nii"
+            # Call applywarp in chunks of volumes
+            warp_in_chunks(
+                name="applywarp",
+                input_path=input["path"].resolve(),
+                output_path=applywarp_out_path.resolve(),
+                cmd=lambda in_path, out_path: [
+                    "FSLOUTPUTTYPE=NIFTI",
+                    "applywarp",
+                    "--interp=spline",
+                    f"-i {in_path}",
+                    # use resampled reference
+                    f"-r {flirt_out_path.resolve()}",
+                    f"-w {warp_file_path.resolve()}",
+                    f"-o {out_path}",
+                ],
+                tempdir=element_tempdir,
+            )
 
             logger.debug("Updating warped data")
             input.update(
@@ -139,11 +147,10 @@ class FSLWarper:
             # Check for data type's mask and warp if found
             if input.get("mask") is not None:
                 # Create a tempfile for warped mask output
-                applywarp_mask_out_path = (
-                    element_tempdir / "warped_mask.nii.gz"
-                )
+                applywarp_mask_out_path = element_tempdir / "warped_mask.nii"
                 # Set applywarp command
                 applywarp_mask_cmd = [
+                    "FSLOUTPUTTYPE=NIFTI",
                     "applywarp",
                     "--interp=nn",
                     f"-i {input['mask']['path'].resolve()}",
@@ -200,10 +207,11 @@ class FSLWarper:
                 logger.debug("Resampling reference to match input resolution")
                 # Create a tempfile for resampled reference output
                 ref_path = (
-                    element_tempdir / f"resampled_reference-{reference}.nii.gz"
+                    element_tempdir / f"resampled_reference-{reference}.nii"
                 )
                 # Set flirt command
                 flirt_cmd = [
+                    "FSLOUTPUTTYPE=NIFTI",
                     "flirt",
                     "-interp spline",
                     f"-in {input['reference']['path'].resolve()}",
@@ -219,21 +227,25 @@ class FSLWarper:
 
             # Create a tempfile for warped output
             applywarp_out_path = (
-                element_tempdir
-                / f"warped_data_from_native_to_{reference}.nii.gz"
+                element_tempdir / f"warped_data_from_native_to_{reference}.nii"
+            )
+            # Call applywarp in chunks of volumes
+            warp_in_chunks(
+                name="applywarp",
+                input_path=input["path"].resolve(),
+                output_path=applywarp_out_path.resolve(),
+                cmd=lambda in_path, out_path: [
+                    "FSLOUTPUTTYPE=NIFTI",
+                    "applywarp",
+                    "--interp=spline",
+                    f"-i {in_path}",
+                    # use resampled reference or original
+                    f"-r {ref_path.resolve()}",
+                    f"-w {warp_file_path.resolve()}",
+                    f"-o {out_path}",
+                ],
+                tempdir=element_tempdir,
             )
-            # Set applywarp command
-            applywarp_cmd = [
-                "applywarp",
-                "--interp=spline",
-                f"-i {input['path'].resolve()}",
-                # use resampled reference or original
-                f"-r {ref_path.resolve()}",
-                f"-w {warp_file_path.resolve()}",
-                f"-o {applywarp_out_path.resolve()}",
-            ]
-            # Call applywarp
-            run_ext_cmd(name="applywarp", cmd=applywarp_cmd)
 
             logger.debug("Updating warped data")
             input.update(
@@ -257,10 +269,11 @@ class FSLWarper:
                 # Create a tempfile for warped mask output
                 applywarp_mask_out_path = (
                     element_tempdir
-                    / f"warped_mask_from_native_to_{reference}.nii.gz"
+                    / f"warped_mask_from_native_to_{reference}.nii"
                 )
                 # Set applywarp command
                 applywarp_mask_cmd = [
+                    "FSLOUTPUTTYPE=NIFTI",
                     "applywarp",
                     "--interp=nn",
                     f"-i {input['mask']['path'].resolve()}",

junifer/preprocess/warping/_utils.py

@@ -0,0 +1,106 @@
+"""Provide utilities for warping."""
+
+# Authors: Federico Raimondo <f.raimondo@fz-juelich.de>
+# License: AGPL
+
+from collections.abc import Callable
+from pathlib import Path
+
+import nibabel as nib
+import numpy as np
+from tqdm import tqdm
+
+from ...utils import config, raise_error, run_ext_cmd
+from ..base import logger
+
+
+__all__ = ["warp_in_chunks"]
+
+
+def warp_in_chunks(
+    name: str,
+    input_path: Path,
+    output_path: Path,
+    cmd: Callable[[Path, Path], list[str]],
+    tempdir: Path,
+) -> None:
+    """Warp an image, in chunks of volumes if it is 4D.
+
+    The memory used by the warping commands grows with the number of volumes,
+    so the volumes of a 4D image are warped in chunks of
+    ``preprocessing.warping.chunksize`` volumes (default 100), which are then
+    concatenated. The volumes are warped independently, so the result is the
+    same as warping the whole image.
+
+    Parameters
+    ----------
+    name : str
+        The name of the warping command.
+    input_path : pathlib.Path
+        The path of the image to warp.
+    output_path : pathlib.Path
+        The path of the warped image.
+    cmd : callable
+        A function that returns the command to warp an image, given the paths
+        of the image to warp and of the warped image.
+    tempdir : pathlib.Path
+        The directory for the chunks.
+
+    Raises
+    ------
+    ValueError
+        If ``preprocessing.warping.chunksize`` is not a positive integer.
+
+    """
+    chunk_size = config.get("preprocessing.warping.chunksize", 100)
+    if (
+        isinstance(chunk_size, bool)
+        or not isinstance(chunk_size, int)
+        or chunk_size < 1
+    ):
+        raise_error(
+            "`preprocessing.warping.chunksize` must be a positive integer, "
+            f"but is {chunk_size}"
+        )
+    img = nib.load(input_path)
+    if img.ndim != 4 or img.shape[3] <= chunk_size:
+        run_ext_cmd(name=name, cmd=cmd(input_path, output_path))
+        return
+
+    n_volumes = img.shape[3]
+    starts = range(0, n_volumes, chunk_size)
+    logger.info(
+        f"Warping {n_volumes} volumes in {len(starts)} chunks of up to "
+        f"{chunk_size} volumes"
+    )
+    warped_chunks = []
+    for i, start in enumerate(tqdm(starts, desc=f"{name} chunks")):
+        # Uncompressed, as they are deleted right after being read
+        chunk_path = tempdir / f"chunk{i}_input_{output_path.stem}.nii"
+        warped_chunk_path = tempdir / f"chunk{i}_{output_path.stem}.nii"
+        chunk = img.slicer[..., start : start + chunk_size]
+        # The data is read scaled, so save it as float instead of scaling it
+        # again (with a different scaling) in its data type
+        if img.dataobj.slope != 1 or img.dataobj.inter != 0:
+            chunk.set_data_dtype(np.float32)
+        nib.save(chunk, chunk_path)
+        run_ext_cmd(name=name, cmd=cmd(chunk_path, warped_chunk_path))
+        chunk_path.unlink()
+        warped_chunks.append(warped_chunk_path)
+
+    # Fill the data of the warped image with the warped chunks (instead of
+    # concatenating them, which needs twice the memory)
+    data = None
+    for start, warped_chunk_path in zip(starts, warped_chunks, strict=True):
+        warped_chunk = nib.load(warped_chunk_path)
+        chunk_data = np.asanyarray(warped_chunk.dataobj)
+        if data is None:
+            first = warped_chunk
+            data = np.empty(
+                (*chunk_data.shape[:3], n_volumes), dtype=chunk_data.dtype
+            )
+        data[..., start : start + chunk_data.shape[3]] = chunk_data
+        warped_chunk_path.unlink()
+    warped = nib.Nifti1Image(data, first.affine, first.header)
+    warped.set_data_dtype(data.dtype)
+    nib.save(warped, output_path)