From 44f75486e13e77584301cf31326c61ab9bd3d2a6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Cl=C3=A9ment=20POIRET?= <poiret.clement@outlook.fr>
Date: Wed, 1 Sep 2021 16:16:12 +0200
Subject: [PATCH] Condition on decrop is zero

---
 pyproject.toml    |  2 +-
 roiloc/locator.py | 15 +++++++++++----
 2 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index be56acb..54b5aac 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "ROILoc"
-version = "0.2.2"
+version = "0.2.3"
 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"
diff --git a/roiloc/locator.py b/roiloc/locator.py
index 057d8dd..d5ef451 100644
--- a/roiloc/locator.py
+++ b/roiloc/locator.py
@@ -120,19 +120,26 @@ def fit_transform(self, image: ANTsImage) -> list:
         self.fit(image)
         return self.transform(image)
 
-    def inverse_transform(self, image: ANTsImage) -> ANTsImage:
+    def inverse_transform(self,
+                          image: ANTsImage,
+                          decrop_is_zero: bool = True) -> ANTsImage:
         """Inverse transform the image to the native space.
 
         Args:
             image (ANTsImage): Image to inverse transform.
+            decrop_is_zero (bool, optional): Whether to decrop the image with
+                zeros or original tissue. Defaults to True.
 
         Returns:
             ANTsImage: Inverse transformed image.
         """
-        empty_image = self._image.new_image_like(
-            np.zeros_like(self._image.numpy()))
+        if decrop_is_zero:
+            reference = self._image.new_image_like(
+                np.zeros_like(self._image.numpy()))
+        else:
+            reference = self._image
 
-        return ants.decrop_image(image, empty_image)
+        return ants.decrop_image(image, reference)
 
 
 def test():