checked warp saving to ants format

fireants/registration/greedy.py

@@ -5,6 +5,7 @@
 import numpy as np
 from typing import List, Optional, Union, Callable
 from tqdm import tqdm
+import SimpleITK as sitk
 
 from fireants.utils.globals import MIN_IMG_SIZE
 from fireants.io.image import BatchedImages
@@ -97,6 +98,39 @@ def get_warped_coordinates(self, fixed_images: BatchedImages, moving_images: Bat
         # move these coordinates, and return them
         moved_coords = fixed_image_affinecoords + warp_field  # affine transform + warp field   
         return moved_coords
+
+    def save_as_ants_transforms(self, filenames: Union[str, List[str]]):
+        ''' given a list of filenames, save the warp field as ants transforms '''
+        if isinstance(filenames, str):
+            filenames = [filenames]
+        assert len(filenames) == self.opt_size, "Number of filenames should match the number of warps"
+        # get the warp field
+        fixed_image: BatchedImages = self.fixed_images
+        moving_image: BatchedImages = self.moving_images
+        # get the moved coordinates
+        moved_coords = self.get_warped_coordinates(fixed_image, moving_image)   # [B, H, W, [D], dim]
+        init_grid = F.affine_grid(torch.eye(self.dims, self.dims+1, device=moved_coords.device)[None], \
+                                  fixed_image.shape, align_corners=True)
+        # this is now moved displacements
+        moved_coords = moved_coords - init_grid
+
+        # convert this grid into moving coordinates 
+        moving_t2p = moving_image.get_torch2phy()[:, :self.dims, :self.dims]
+        moved_coords = torch.einsum('bij, b...j->b...i', moving_t2p, moved_coords)
+        # save 
+        for i in range(self.opt_size):
+            moved_disp = moved_coords[i].detach().cpu().numpy()  # [H, W, D, 3]
+            savefile = filenames[i]
+            # get itk image
+            if len(fixed_image.images) < i:     # this image is probably broadcasted then
+                itk_data = fixed_image.images[0].itk_image
+            else:
+                itk_data = fixed_image.images[i].itk_image
+            # copy itk data
+            warp = sitk.GetImageFromArray(moved_disp)
+            warp.CopyInformation(itk_data)
+            sitk.WriteImage(warp, savefile)
+
 
     def optimize(self, save_transformed=False):
         ''' optimize the warp field to match the two images based on loss function '''