From 58d29bf8d97c75e064f4212fa10856f30b6acbef Mon Sep 17 00:00:00 2001
From: ncullen93 <nicholas.cullen@med.lu.se>
Date: Thu, 30 May 2024 11:31:54 +0200
Subject: [PATCH] ENH: crop vector images

---
 ants/core/ants_image_io.py | 16 ++++++++--------
 ants/ops/crop_image.py     |  9 +++++++++
 ants/ops/reorient_image.py |  3 ++-
 ants/ops/resample_image.py |  2 +-
 tests/test_utils.py        |  8 ++++++++
 5 files changed, 28 insertions(+), 10 deletions(-)

diff --git a/ants/core/ants_image_io.py b/ants/core/ants_image_io.py
index f0c21dff..88210318 100644
--- a/ants/core/ants_image_io.py
+++ b/ants/core/ants_image_io.py
@@ -453,7 +453,7 @@ def image_write(image, filename, ri=False):
         return image
 
 @image_method
-def clone(self, pixeltype=None):
+def clone(image, pixeltype=None):
     """
     Create a copy of the given ANTsImage with the same data and info, possibly with
     a different data type for the image data. Only supports casting to
@@ -476,22 +476,22 @@ def clone(self, pixeltype=None):
     ANTsImage
     """
     if pixeltype is None:
-        pixeltype = self.pixeltype
+        pixeltype = image.pixeltype
 
     if pixeltype not in _supported_ptypes:
         raise ValueError('Pixeltype %s not supported. Supported types are %s' % (pixeltype, _supported_ptypes))
 
-    if self.has_components and (not self.is_rgb):
-        comp_imgs = ants.split_channels(self)
+    if image.has_components and (not image.is_rgb):
+        comp_imgs = ants.split_channels(image)
         comp_imgs_cloned = [comp_img.clone(pixeltype) for comp_img in comp_imgs]
-        return ants.merge_channels(comp_imgs_cloned)
+        return ants.merge_channels(comp_imgs_cloned, channels_first=image.channels_first)
     else:
-        p1_short = short_ptype(self.pixeltype)
+        p1_short = short_ptype(image.pixeltype)
         p2_short = short_ptype(pixeltype)
-        ndim = self.dimension
+        ndim = image.dimension
         fn_suffix = '%s%i' % (p2_short,ndim)
         libfn = get_lib_fn('antsImageClone%s'%fn_suffix)
-        pointer_cloned = libfn(self.pointer)
+        pointer_cloned = libfn(image.pointer)
         return ants.from_pointer(pointer_cloned)
         
 copy = clone
diff --git a/ants/ops/crop_image.py b/ants/ops/crop_image.py
index 8458b9ea..9120fbc1 100644
--- a/ants/ops/crop_image.py
+++ b/ants/ops/crop_image.py
@@ -37,8 +37,14 @@ def crop_image(image, label_image=None, label=1):
     >>> import ants
     >>> fi = ants.image_read( ants.get_ants_data('r16') )
     >>> cropped = ants.crop_image(fi)
+    >>> fi2 = ants.merge_channels([fi,fi])
+    >>> cropped2 = ants.crop_image(fi2)
     >>> cropped = ants.crop_image(fi, fi, 100 )
     """
+    if image.has_components:
+        return ants.merge_channels([crop_image(img, label_image, label) for img in ants.split_channels(image)],
+                                   channels_first=image.channels_first)
+        
     inpixeltype = image.pixeltype
     ndim = image.dimension
     if image.pixeltype != 'float':
@@ -88,6 +94,9 @@ def crop_indices(image, lowerind, upperind):
     >>> cropped = ants.smooth_image( cropped, 5 )
     >>> decropped = ants.decrop_image( cropped, fi )
     """
+    if image.has_components:
+        return ants.merge_channels([crop_indices(img, lowerind, upperind) for img in ants.split_channels(image)])
+        
     inpixeltype = 'float'
     if image.pixeltype != 'float':
         inpixeltype = image.pixeltype
diff --git a/ants/ops/reorient_image.py b/ants/ops/reorient_image.py
index 973dcc0e..22f81232 100644
--- a/ants/ops/reorient_image.py
+++ b/ants/ops/reorient_image.py
@@ -63,7 +63,8 @@ def reorient_image2(image, orientation='RAS'):
      >>> mni2 = mni.reorient_image2()
      """
      if image.has_components:
-         return ants.merge_channels([img.reorient_image2(orientation) for img in ants.split_channels(image)])
+         return ants.merge_channels([img.reorient_image2(orientation) for img in ants.split_channels(image)],
+                                    channels_first=image.channels_first)
      
      if image.dimension != 3:
          raise ValueError('image must have 3 dimensions')
diff --git a/ants/ops/resample_image.py b/ants/ops/resample_image.py
index 16ae8861..5bc19631 100644
--- a/ants/ops/resample_image.py
+++ b/ants/ops/resample_image.py
@@ -72,7 +72,7 @@ def resample_image(image, resample_params, use_voxels=False, interp_type=1):
             libfn(processed_args)
             outimage = outimage.clone(image.pixeltype)
             new_images.append(outimage)
-        outimage = ants.merge_channels(new_images)
+        outimage = ants.merge_channels(new_images, channels_first=image.channels_first)
         return outimage
 
 @image_method
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 18e691af..edcfc132 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -209,6 +209,14 @@ def test_crop_image_example(self):
 
         # label image not float
         cropped = ants.crop_image(fi, fi.clone("unsigned int"), 100)
+        
+        # channel image
+        fi = ants.image_read( ants.get_ants_data('r16') )
+        cropped = ants.crop_image(fi)
+        fi2 = ants.merge_channels([fi,fi])
+        cropped2 = ants.crop_image(fi2)
+        
+        self.assertEqual(cropped.shape, cropped2.shape)
 
     def test_crop_indices_example(self):
         fi = ants.image_read(ants.get_ants_data("r16"))