New RoiLocator class

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "ROILoc"
-version = "0.1.3"
+version = "0.2.0"
 description = "A simple package to center and crop T1w & T2w MRIs around a given region of interest by its name."
 license = "MIT"
 readme = "README.rst"

roiloc/__init__.py

@@ -1 +1 @@
-__version__ = "0.1.3"
+__version__ = "0.2.0"

roiloc/location.py

@@ -1,4 +1,5 @@
 from pathlib import PosixPath
+from typing import Optional
 
 import ants
 import numpy as np
@@ -36,27 +37,32 @@ def get_coords(x: np.ndarray, margin: list = [4, 4, 2]) -> list:
 
 def crop(image: ANTsImage,
          coords: list,
-         output_path: PosixPath,
-         log_coords: bool = True):
+         output_path: Optional[PosixPath] = None,
+         log_coords: bool = True,
+         ri: bool = False):
     """Crop an image using coordinates.
 
     Args:
         image (ANTsImage): image to be cropped
         coords (list): coordinates of the ROI
-        output_path (PosixPath): path to save the cropped image
+        output_path (PosixPath, optional): path to save the cropped image
         log_coords (bool, optional): log the coordinates. Defaults to True.
+        ri (bool): if True, return the ROI as an ANTsImage. Defaults to False.
     """
     cropped_image = ants.crop_indices(image,
                                       lowerind=coords[:3],
                                       upperind=coords[3:])
 
     if cropped_image.numpy().any():
-        ants.image_write(cropped_image, str(output_path), ri=False)
-
-        if log_coords:
-            np.savetxt(output_path.with_suffix(".txt"), coords)
+        if output_path:
+            ants.image_write(cropped_image, str(output_path), ri=False)
+            if log_coords:
+                np.savetxt(output_path.with_suffix(".txt"), coords)
 
     else:
         print(
-            f"[italic white]\tEmpty cropped array, skipping {output_path} for coordinates {coords}..."
+            f"[italic white]\tEmpty cropped array, skipping for coordinates {coords}..."
         )
+
+    if ri:
+        return cropped_image

roiloc/locator.py

@@ -0,0 +1,85 @@
+from typing import Optional
+
+import ants
+from ants.core import ANTsImage
+
+from .location import crop, get_coords
+from .registration import get_roi, register
+from .template import get_atlas, get_mni, get_roi_indices
+
+
+class RoiLocator:
+
+    def __init__(self,
+                 contrast: str,
+                 roi: str,
+                 bet: bool = False,
+                 transform_type: str = "AffineFast",
+                 margin: list = [4, 4, 4],
+                 mask: Optional[ANTsImage] = None):
+        self.transform_type = transform_type
+        self.margin = margin
+        self.mask = mask
+
+        self._roi_idx = get_roi_indices(roi)
+        self._mni = get_mni(contrast, bet)
+        self._atlas = get_atlas()
+
+        self._image = None
+        self._fwdtransforms = None
+        self._invtransforms = None
+        self.coords = {}
+
+    def get_coords(self) -> dict:
+        return self.coords
+
+    def fit(self, image: ANTsImage):
+        self._image = image
+
+        registration = ants.registration(fixed=image,
+                                         moving=self._mni,
+                                         type_of_transform=self.transform_type,
+                                         mask=self.mask)
+
+        self._fwdtransforms = registration["fwdtransforms"]
+        self._invtransforms = registration["invtransforms"]
+
+        registered_atlas = ants.apply_transforms(
+            fixed=image,
+            moving=self._atlas,
+            transformlist=self._fwdtransforms,
+            interpolator="nearestNeighbor")
+
+        for i, side in enumerate(["right", "left"]):
+            region = get_roi(registered_atlas=registered_atlas,
+                             idx=int(self._roi_idx[i]),
+                             save=False)
+            coords = get_coords(region.numpy(), margin=self.margin)
+
+            self.coords[side] = coords
+
+    def transform(self, image: ANTsImage) -> list:
+        return [
+            crop(image, self.coords[side], log_coords=False, ri=True)
+            for side in ["right", "left"]
+        ]
+
+    def fit_transform(self, image: ANTsImage) -> list:
+        self.fit(image)
+        return self.transform(image)
+
+    def inverse_transform(self, image: ANTsImage) -> ANTsImage:
+        return ants.decrop_image(image, self._image)
+
+
+def test():
+    image = ants.image_read(
+        "~/Datasets/MemoDev/ManualSegmentation/mb190108/tse.nii.gz",
+        reorient="LPI")
+    locator = RoiLocator(contrast="t2", roi="hippocampus", bet=True)
+
+    right, _ = locator.fit_transform(image)
+    print(locator.get_coords())
+    right_orig = locator.inverse_transform(right)
+
+    assert right_orig.shape == image.shape

roiloc/registration.py

@@ -43,8 +43,8 @@ def register(fixed: ANTsImage,
 
 def get_roi(registered_atlas: ANTsImage,
             idx: int,
-            output_dir: str,
-            output_file: str,
+            output_dir: Optional[str] = None,
+            output_file: Optional[str] = None,
             save: bool = True) -> ANTsImage:
     """Get the registered ROI from CerebrA atlas, into a
     subject's native space.
@@ -54,8 +54,8 @@ def get_roi(registered_atlas: ANTsImage,
         atlas (ANTsImage): CerebrA Atlas
         idx (int): Index of the ROI
         transform (list): Transformation from MNI to Native space
-        output_dir (str): Where to save the ROIs
-        output_file (str): Name of the ROIs
+        output_dir (str, optional): Where to save the ROIs
+        output_file (str, optional): Name of the ROIs
         save (bool, optional): Save or not the ROIs. Defaults to True.
 
     Returns:

roiloc/sanitychecks.py

@@ -1,13 +1,8 @@
-from typing import Union
-
 import ants
-import nibabel as nib
 from ants.core.ants_image import ANTsImage
-from nibabel.nifti1 import Nifti1Image
-from nibabel.nifti2 import Nifti2Image
 
 
-def is_lpi_ants(image: ANTsImage) -> bool:
+def is_lpi(image: ANTsImage) -> bool:
     """Check if the ants image is LPI-
 
     Args:
@@ -17,15 +12,3 @@ def is_lpi_ants(image: ANTsImage) -> bool:
         bool: LPI- status
     """
     return ants.get_orientation(image) == "LPI"
-
-
-def is_ras_nib(image: Union[Nifti1Image, Nifti2Image]) -> bool:
-    """Check if the ants image is RAS+
-
-    Args:
-        image (Union[Nifti1Image, Nifti2Image]): MRI
-
-    Returns:
-        bool: RAS+ status
-    """
-    return "".join(nib.aff2axcodes(image.affine)) == "RAS"

tests/test_roiloc.py

@@ -2,4 +2,4 @@
 
 
 def test_version():
-    assert __version__ == '0.1.3'
+    assert __version__ == '0.2.0'