Set proper affine when cropping patches

fepegar · Jul 1, 2020 · a6d895e · a6d895e
1 parent 0ec0360
commit a6d895e
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/torchio/data/image.py b/torchio/data/image.py
@@ -298,8 +298,10 @@ def set_check_nans(self, check_nans):
         self.check_nans = check_nans
 
     def crop(self, index_ini, index_fin):
-        # TODO: handle affine
+        new_origin = nib.affines.apply_affine(self.affine, index_ini)
+        new_affine = self.affine.copy()
+        new_affine[:3, 3] = new_origin
         i0, j0, k0 = index_ini
         i1, j1, k1 = index_fin
         patch = self.data[0, i0:i1, j0:j1, k0:k1].clone()
-        return Image(tensor=patch, affine=self.affine, type=self.type)
+        return Image(tensor=patch, affine=new_affine, type=self.type)