Stop loading data when copying unloaded image (#982)

* Add failing test * Stop loading data when copying image * Add assertions after loading as well Co-authored-by: Derk <derk@mriguidance.com>
TorchIO-project · Oct 11, 2022 · 134fa4a · 134fa4a
1 parent 74457f8
commit 134fa4a
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 4 deletions.
diff --git a/src/torchio/constants.py b/src/torchio/constants.py
@@ -11,6 +11,7 @@
 STEM = 'stem'
 DATA = 'data'
 AFFINE = 'affine'
+TENSOR = 'tensor'
 
 # For aggregator
 IMAGE = 'image'

diff --git a/src/torchio/data/image.py b/src/torchio/data/image.py
@@ -23,6 +23,7 @@
 from ..constants import LABEL
 from ..constants import PATH
 from ..constants import STEM
+from ..constants import TENSOR
 from ..constants import TYPE
 from ..typing import TypeData
 from ..typing import TypeDataAffine
@@ -204,11 +205,12 @@ def __array__(self):
 
     def __copy__(self):
         kwargs = {
-            'tensor': self.data,
-            'affine': self.affine,
-            'type': self.type,
-            'path': self.path,
+            TYPE: self.type,
+            PATH: self.path,
         }
+        if self._loaded:
+            kwargs[TENSOR] = self.data
+            kwargs[AFFINE] = self.affine
         for key, value in self.items():
             if key in PROTECTED_KEYS:
                 continue

diff --git a/tests/data/test_image.py b/tests/data/test_image.py
@@ -247,3 +247,17 @@ def numpy_reader(path):
         np.save(test_path, tensor)
         image = tio.ScalarImage(test_path, reader=numpy_reader)
         image.load()
+
+    def test_copy_no_data(self):
+        # https://github.com/fepegar/torchio/issues/974
+        path = self.get_image_path('im_copy')
+        my_image = tio.LabelMap(path)
+        assert not my_image._loaded
+        new_image = copy.copy(my_image)
+        assert not my_image._loaded
+        assert not new_image._loaded
+
+        my_image.load()
+        new_image = copy.copy(my_image)
+        assert my_image._loaded
+        assert new_image._loaded