junifer 0.0.5.dev242__py3-none-any.whl → 0.0.6__py3-none-any.whl

junifer/preprocess/warping/_ants_warper.py

@@ -6,10 +6,6 @@
 from typing import (
     Any,
     ClassVar,
-    Dict,
-    List,
-    Set,
-    Union,
 )
 
 import nibabel as nib
@@ -17,7 +13,8 @@ import numpy as np
 
 from ...data import get_template, get_xfm
 from ...pipeline import WorkDirManager
-from ...utils import logger, run_ext_cmd
+from ...typing import Dependencies, ExternalDependencies
+from ...utils import logger, raise_error, run_ext_cmd
 
 
 __all__ = ["ANTsWarper"]
@@ -31,21 +28,21 @@ class ANTsWarper:
 
     """
 
-    _EXT_DEPENDENCIES: ClassVar[List[Dict[str, Union[str, List[str]]]]] = [
+    _EXT_DEPENDENCIES: ClassVar[ExternalDependencies] = [
         {
             "name": "ants",
             "commands": ["ResampleImage", "antsApplyTransforms"],
         },
     ]
 
-    _DEPENDENCIES: ClassVar[Set[str]] = {"numpy", "nibabel"}
+    _DEPENDENCIES: ClassVar[Dependencies] = {"numpy", "nibabel"}
 
     def preprocess(
         self,
-        input: Dict[str, Any],
-        extra_input: Dict[str, Any],
+        input: dict[str, Any],
+        extra_input: dict[str, Any],
         reference: str,
-    ) -> Dict[str, Any]:
+    ) -> dict[str, Any]:
         """Preprocess using ANTs.
 
         Parameters
@@ -61,9 +58,12 @@ class ANTsWarper:
         Returns
         -------
         dict
-            The ``input`` dictionary with modified ``data`` and ``space`` key
-            values and new ``reference_path`` key whose value points to the
-            reference file used for warping.
+            The ``input`` dictionary with updated values.
+
+        Raises
+        ------
+        RuntimeError
+            If warp file path could not be found in ``extra_input``.
 
         """
         # Create element-specific tempdir for storing post-warping assets
@@ -79,6 +79,17 @@ class ANTsWarper:
             # resolution
             resolution = np.min(input["data"].header.get_zooms()[:3])
 
+            # Get warp file path
+            warp_file_path = None
+            for entry in extra_input["Warp"]:
+                if entry["dst"] == "native":
+                    warp_file_path = entry["path"]
+            if warp_file_path is None:
+                raise_error(
+                    klass=RuntimeError,
+                    msg="Could not find correct warp file path",
+                )
+
             # Create a tempfile for resampled reference output
             resample_image_out_path = (
                 element_tempdir / "resampled_reference.nii.gz"
@@ -97,7 +108,7 @@ class ANTsWarper:
             run_ext_cmd(name="ResampleImage", cmd=resample_image_cmd)
 
             # Create a tempfile for warped output
-            apply_transforms_out_path = element_tempdir / "output.nii.gz"
+            apply_transforms_out_path = element_tempdir / "warped_data.nii.gz"
             # Set antsApplyTransforms command
             apply_transforms_cmd = [
                 "antsApplyTransforms",
@@ -107,18 +118,65 @@ class ANTsWarper:
                 f"-i {input['path'].resolve()}",
                 # use resampled reference
                 f"-r {resample_image_out_path.resolve()}",
-                f"-t {extra_input['Warp']['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)
 
-            # Load nifti
-            input["data"] = nib.load(apply_transforms_out_path)
-            # Save resampled reference path
-            input["reference_path"] = resample_image_out_path
-            # Use reference input's space as warped input's space
-            input["space"] = extra_input["T1w"]["space"]
+            logger.debug("Updating warped data")
+            input.update(
+                {
+                    # Update path to sync with "data"
+                    "path": apply_transforms_out_path,
+                    # Load nifti
+                    "data": nib.load(apply_transforms_out_path),
+                    # Use reference input's space as warped input's space
+                    "space": extra_input["T1w"]["space"],
+                    # Save resampled reference path
+                    "reference": {"path": resample_image_out_path},
+                    # Keep pre-warp space for further operations
+                    "prewarp_space": input["space"],
+                }
+            )
+
+            # Check for data type's mask and warp if found
+            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"
+                )
+                # Set antsApplyTransforms command
+                apply_transforms_mask_cmd = [
+                    "antsApplyTransforms",
+                    "-d 3",
+                    "-e 3",
+                    "-n 'GenericLabel[NearestNeighbor]'",
+                    f"-i {input['mask']['path'].resolve()}",
+                    # use resampled reference
+                    f"-r {input['reference']['path'].resolve()}",
+                    f"-t {warp_file_path.resolve()}",
+                    f"-o {apply_transforms_mask_out_path.resolve()}",
+                ]
+                # Call antsApplyTransforms
+                run_ext_cmd(
+                    name="antsApplyTransforms", cmd=apply_transforms_mask_cmd
+                )
+
+                logger.debug("Updating warped mask data")
+                input.update(
+                    {
+                        "mask": {
+                            # Update path to sync with "data"
+                            "path": apply_transforms_mask_out_path,
+                            # Load nifti
+                            "data": nib.load(apply_transforms_mask_out_path),
+                            # Use reference input's space as warped input
+                            # mask's space
+                            "space": extra_input["T1w"]["space"],
+                        }
+                    }
+                )
 
         # Template space warping
         else:
@@ -131,19 +189,18 @@ class ANTsWarper:
             # Get template space image
             template_space_img = get_template(
                 space=reference,
-                target_data=input,
+                target_img=input["data"],
                 extra_input=None,
             )
-
-            # Create component-scoped tempdir
-            tempdir = WorkDirManager().get_tempdir(prefix="ants_warper")
             # Save template
-            template_space_img_path = tempdir / f"{reference}_T1w.nii.gz"
+            template_space_img_path = (
+                element_tempdir / f"{reference}_T1w.nii.gz"
+            )
             nib.save(template_space_img, template_space_img_path)
 
             # Create a tempfile for warped output
             warped_output_path = element_tempdir / (
-                f"data_warped_from_{input['space']}_to_" f"{reference}.nii.gz"
+                f"warped_data_from_{input['space']}_to_{reference}.nii.gz"
             )
 
             # Set antsApplyTransforms command
@@ -160,11 +217,58 @@ class ANTsWarper:
             # Call antsApplyTransforms
             run_ext_cmd(name="antsApplyTransforms", cmd=apply_transforms_cmd)
 
-            # Delete tempdir
-            WorkDirManager().delete_tempdir(tempdir)
+            logger.debug("Updating warped data")
+            input.update(
+                {
+                    # Update path to sync with "data"
+                    "path": warped_output_path,
+                    # Load nifti
+                    "data": nib.load(warped_output_path),
+                    # Update warped input's space
+                    "space": reference,
+                    # Save reference path
+                    "reference": {"path": template_space_img_path},
+                    # Keep pre-warp space for further operations
+                    "prewarp_space": input["space"],
+                }
+            )
+
+            # Check for data type's mask and warp if found
+            if input.get("mask") is not None:
+                # Create a tempfile for warped mask output
+                apply_transforms_mask_out_path = element_tempdir / (
+                    f"warped_mask_from_{input['space']}_to_"
+                    f"{reference}.nii.gz"
+                )
+                # Set antsApplyTransforms command
+                apply_transforms_mask_cmd = [
+                    "antsApplyTransforms",
+                    "-d 3",
+                    "-e 3",
+                    "-n 'GenericLabel[NearestNeighbor]'",
+                    f"-i {input['mask']['path'].resolve()}",
+                    # use resampled reference
+                    f"-r {input['reference']['path'].resolve()}",
+                    f"-t {xfm_file_path.resolve()}",
+                    f"-o {apply_transforms_mask_out_path.resolve()}",
+                ]
+                # Call antsApplyTransforms
+                run_ext_cmd(
+                    name="antsApplyTransforms", cmd=apply_transforms_mask_cmd
+                )
 
-            # Modify target data
-            input["data"] = nib.load(warped_output_path)
-            input["space"] = reference
+                logger.debug("Updating warped mask data")
+                input.update(
+                    {
+                        "mask": {
+                            # Update path to sync with "data"
+                            "path": apply_transforms_mask_out_path,
+                            # Load nifti
+                            "data": nib.load(apply_transforms_mask_out_path),
+                            # Update warped input mask's space
+                            "space": reference,
+                        }
+                    }
+                )
 
         return input

junifer/preprocess/warping/_fsl_warper.py

@@ -6,17 +6,14 @@
 from typing import (
     Any,
     ClassVar,
-    Dict,
-    List,
-    Set,
-    Union,
 )
 
 import nibabel as nib
 import numpy as np
 
 from ...pipeline import WorkDirManager
-from ...utils import logger, run_ext_cmd
+from ...typing import Dependencies, ExternalDependencies
+from ...utils import logger, raise_error, run_ext_cmd
 
 
 __all__ = ["FSLWarper"]
@@ -30,20 +27,20 @@ class FSLWarper:
 
     """
 
-    _EXT_DEPENDENCIES: ClassVar[List[Dict[str, Union[str, List[str]]]]] = [
+    _EXT_DEPENDENCIES: ClassVar[ExternalDependencies] = [
         {
             "name": "fsl",
             "commands": ["flirt", "applywarp"],
         },
     ]
 
-    _DEPENDENCIES: ClassVar[Set[str]] = {"numpy", "nibabel"}
+    _DEPENDENCIES: ClassVar[Dependencies] = {"numpy", "nibabel"}
 
     def preprocess(
         self,
-        input: Dict[str, Any],
-        extra_input: Dict[str, Any],
-    ) -> Dict[str, Any]:
+        input: dict[str, Any],
+        extra_input: dict[str, Any],
+    ) -> dict[str, Any]:
         """Preprocess using FSL.
 
         Parameters
@@ -57,9 +54,12 @@ class FSLWarper:
         Returns
         -------
         dict
-            The ``input`` dictionary with modified ``data`` and ``space`` key
-            values and new ``reference_path`` key whose value points to the
-            reference file used for warping.
+            The ``input`` dictionary with updated values.
+
+        Raises
+        ------
+        RuntimeError
+            If warp file path could not be found in ``extra_input``.
 
         """
         logger.debug("Using FSL for space warping")
@@ -68,6 +68,16 @@ class FSLWarper:
         # resolution
         resolution = np.min(input["data"].header.get_zooms()[:3])
 
+        # Get warp file path
+        warp_file_path = None
+        for entry in extra_input["Warp"]:
+            if entry["dst"] == "native":
+                warp_file_path = entry["path"]
+        if warp_file_path is None:
+            raise_error(
+                klass=RuntimeError, msg="Could not find correct warp file path"
+            )
+
         # Create element-specific tempdir for storing post-warping assets
         element_tempdir = WorkDirManager().get_element_tempdir(
             prefix="fsl_warper"
@@ -88,25 +98,66 @@ class FSLWarper:
         run_ext_cmd(name="flirt", cmd=flirt_cmd)
 
         # Create a tempfile for warped output
-        applywarp_out_path = element_tempdir / "output.nii.gz"
+        applywarp_out_path = element_tempdir / "warped_data.nii.gz"
         # Set applywarp command
         applywarp_cmd = [
             "applywarp",
             "--interp=spline",
             f"-i {input['path'].resolve()}",
-            f"-r {flirt_out_path.resolve()}",  # use resampled reference
-            f"-w {extra_input['Warp']['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)
 
-        # Load nifti
-        input["data"] = nib.load(applywarp_out_path)
-        # Save resampled reference path
-        input["reference_path"] = flirt_out_path
+        logger.debug("Updating warped data")
+        input.update(
+            {
+                # Update path to sync with "data"
+                "path": applywarp_out_path,
+                # Load nifti
+                "data": nib.load(applywarp_out_path),
+                # Use reference input's space as warped input's space
+                "space": extra_input["T1w"]["space"],
+                # Save resampled reference path
+                "reference": {"path": flirt_out_path},
+                # Keep pre-warp space for further operations
+                "prewarp_space": input["space"],
+            }
+        )
 
-        # Use reference input's space as warped input's space
-        input["space"] = extra_input["T1w"]["space"]
+        # 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"
+            # Set applywarp command
+            applywarp_mask_cmd = [
+                "applywarp",
+                "--interp=nn",
+                f"-i {input['mask']['path'].resolve()}",
+                # use resampled reference
+                f"-r {input['reference']['path'].resolve()}",
+                f"-w {warp_file_path.resolve()}",
+                f"-o {applywarp_mask_out_path.resolve()}",
+            ]
+            # Call applywarp
+            run_ext_cmd(name="applywarp", cmd=applywarp_mask_cmd)
+
+            logger.debug("Updating warped mask data")
+            input.update(
+                {
+                    "mask": {
+                        # Update path to sync with "data"
+                        "path": applywarp_mask_out_path,
+                        # Load nifti
+                        "data": nib.load(applywarp_mask_out_path),
+                        # Use reference input's space as warped input mask's
+                        # space
+                        "space": extra_input["T1w"]["space"],
+                    }
+                }
+            )
 
         return input

junifer/preprocess/warping/space_warper.py

@@ -3,11 +3,12 @@
 # Authors: Synchon Mandal <s.mandal@fz-juelich.de>
 # License: AGPL
 
-from typing import Any, ClassVar, Dict, List, Optional, Tuple, Type, Union
+from typing import Any, ClassVar, Optional, Union
 
 from templateflow import api as tflow
 
 from ...api.decorators import register_preprocessor
+from ...typing import ConditionalDependencies
 from ...utils import logger, raise_error
 from ..base import BasePreprocessor
 from ._ants_warper import ANTsWarper
@@ -23,11 +24,12 @@ class SpaceWarper(BasePreprocessor):
 
     Parameters
     ----------
-    using : {"fsl", "ants"}
+    using : {"fsl", "ants", "auto"}
         Implementation to use for warping:
 
         * "fsl" : Use FSL's ``applywarp``
         * "ants" : Use ANTs' ``antsApplyTransforms``
+        * "auto" : Auto-select tool when ``reference="T1w"``
 
     reference : str
         The data type to use as reference for warping, can be either a data
@@ -46,7 +48,7 @@ class SpaceWarper(BasePreprocessor):
 
     """
 
-    _CONDITIONAL_DEPENDENCIES: ClassVar[List[Dict[str, Union[str, Type]]]] = [
+    _CONDITIONAL_DEPENDENCIES: ClassVar[ConditionalDependencies] = [
         {
             "using": "fsl",
             "depends_on": FSLWarper,
@@ -55,10 +57,14 @@ class SpaceWarper(BasePreprocessor):
             "using": "ants",
             "depends_on": ANTsWarper,
         },
+        {
+            "using": "auto",
+            "depends_on": [FSLWarper, ANTsWarper],
+        },
     ]
 
     def __init__(
-        self, using: str, reference: str, on: Union[List[str], str]
+        self, using: str, reference: str, on: Union[list[str], str]
     ) -> None:
         """Initialize the class."""
         # Validate `using` parameter
@@ -88,7 +94,7 @@ class SpaceWarper(BasePreprocessor):
         else:
             raise_error(f"Unknown reference: {self.reference}")
 
-    def get_valid_inputs(self) -> List[str]:
+    def get_valid_inputs(self) -> list[str]:
         """Get valid data types for input.
 
         Returns
@@ -129,9 +135,9 @@ class SpaceWarper(BasePreprocessor):
 
     def preprocess(
         self,
-        input: Dict[str, Any],
-        extra_input: Optional[Dict[str, Any]] = None,
-    ) -> Tuple[Dict[str, Any], Optional[Dict[str, Dict[str, Any]]]]:
+        input: dict[str, Any],
+        extra_input: Optional[dict[str, Any]] = None,
+    ) -> tuple[dict[str, Any], Optional[dict[str, dict[str, Any]]]]:
         """Preprocess.
 
         Parameters
@@ -155,14 +161,16 @@ class SpaceWarper(BasePreprocessor):
             If ``extra_input`` is None when transforming to native space
             i.e., using ``"T1w"`` as reference.
         RuntimeError
-            If the data is in the correct space and does not require
+            If warper could not be found in ``extra_input`` when
+            ``using="auto"`` or
+            if the data is in the correct space and does not require
             warping or
-            if FSL is used for template space warping.
+            if FSL is used when ``reference="T1w"``.
 
         """
         logger.info(f"Warping to {self.reference} space using SpaceWarper")
         # Transform to native space
-        if self.using in ["fsl", "ants"] and self.reference == "T1w":
+        if self.using in ["fsl", "ants", "auto"] and self.reference == "T1w":
             # Check for extra inputs
             if extra_input is None:
                 raise_error(
@@ -181,6 +189,26 @@ class SpaceWarper(BasePreprocessor):
                     extra_input=extra_input,
                     reference=self.reference,
                 )
+            elif self.using == "auto":
+                warper = None
+                for entry in extra_input["Warp"]:
+                    if entry["dst"] == "native":
+                        warper = entry["warper"]
+                if warper is None:
+                    raise_error(
+                        klass=RuntimeError, msg="Could not find correct warper"
+                    )
+                if warper == "fsl":
+                    input = FSLWarper().preprocess(
+                        input=input,
+                        extra_input=extra_input,
+                    )
+                elif warper == "ants":
+                    input = ANTsWarper().preprocess(
+                        input=input,
+                        extra_input=extra_input,
+                        reference=self.reference,
+                    )
         # Transform to template space with ANTs possible
         elif self.using == "ants" and self.reference != "T1w":
             # Check pre-requirements for space manipulation

junifer/preprocess/warping/tests/test_space_warper.py

@@ -4,7 +4,6 @@
 # License: AGPL
 
 import socket
-from typing import TYPE_CHECKING, Tuple, Type
 
 import pytest
 from numpy.testing import assert_array_equal, assert_raises
@@ -14,10 +13,7 @@ from junifer.datareader import DefaultDataReader
 from junifer.pipeline.utils import _check_ants, _check_fsl
 from junifer.preprocess import SpaceWarper
 from junifer.testing.datagrabbers import PartlyCloudyTestingDataGrabber
-
-
-if TYPE_CHECKING:
-    from junifer.datagrabber import BaseDataGrabber
+from junifer.typing import DataGrabberLike
 
 
 @pytest.mark.parametrize(
@@ -32,7 +28,7 @@ if TYPE_CHECKING:
 def test_SpaceWarper_errors(
     using: str,
     reference: str,
-    error_type: Type[Exception],
+    error_type: type[Exception],
     error_msg: str,
 ) -> None:
     """Test SpaceWarper errors.
@@ -99,7 +95,7 @@ def test_SpaceWarper_errors(
     reason="only for juseless",
 )
 def test_SpaceWarper_native(
-    datagrabber: "BaseDataGrabber", element: Tuple[str, ...], using: str
+    datagrabber: DataGrabberLike, element: tuple[str, ...], using: str
 ) -> None:
     """Test SpaceWarper for native space warping.
 
@@ -162,8 +158,8 @@ def test_SpaceWarper_native(
     _check_ants() is False, reason="requires ANTs to be in PATH"
 )
 def test_SpaceWarper_multi_mni(
-    datagrabber: "BaseDataGrabber",
-    element: Tuple[str, ...],
+    datagrabber: DataGrabberLike,
+    element: tuple[str, ...],
     space: str,
 ) -> None:
     """Test SpaceWarper for MNI space warping.

junifer/stats.py

@@ -4,7 +4,7 @@
 #          Synchon Mandal <s.mandal@fz-juelich.de>
 # License: AGPL
 
-from typing import Any, Callable, Dict, List, Optional
+from typing import Any, Callable, Optional
 
 import numpy as np
 from scipy.stats import mode, trim_mean
@@ -13,11 +13,11 @@ from scipy.stats.mstats import winsorize
 from .utils import logger, raise_error
 
 
-__all__ = ["get_aggfunc_by_name", "count", "winsorized_mean", "select"]
+__all__ = ["count", "get_aggfunc_by_name", "select", "winsorized_mean"]
 
 
 def get_aggfunc_by_name(
-    name: str, func_params: Optional[Dict[str, Any]] = None
+    name: str, func_params: Optional[dict[str, Any]] = None
 ) -> Callable:
     """Get an aggregation function by its name.
 
@@ -169,8 +169,8 @@ def winsorized_mean(
 def select(
     data: np.ndarray,
     axis: int = 0,
-    pick: Optional[List[int]] = None,
-    drop: Optional[List[int]] = None,
+    pick: Optional[list[int]] = None,
+    drop: Optional[list[int]] = None,
 ) -> np.ndarray:
     """Select a subset of the data.
 

junifer/storage/__init__.py

@@ -4,15 +4,7 @@
 #          Synchon Mandal <s.mandal@fz-juelich.de>
 # License: AGPL
 
-from .base import BaseFeatureStorage
-from .pandas_base import PandasBaseFeatureStorage
-from .sqlite import SQLiteFeatureStorage
-from .hdf5 import HDF5FeatureStorage
+import lazy_loader as lazy
 
 
-__all__ = [
-    "BaseFeatureStorage",
-    "PandasBaseFeatureStorage",
-    "SQLiteFeatureStorage",
-    "HDF5FeatureStorage",
-]
+__getattr__, __dir__, __all__ = lazy.attach_stub(__name__, __file__)

junifer/storage/__init__.pyi

@@ -0,0 +1,11 @@
+__all__ = [
+    "BaseFeatureStorage",
+    "PandasBaseFeatureStorage",
+    "SQLiteFeatureStorage",
+    "HDF5FeatureStorage",
+]
+
+from .base import BaseFeatureStorage
+from .pandas_base import PandasBaseFeatureStorage
+from .sqlite import SQLiteFeatureStorage
+from .hdf5 import HDF5FeatureStorage