Updating the atlas-based segmentation

clabtoolkit/imagetools.py

@@ -174,33 +174,43 @@ def cropped_to_native(in_image: str, native_image: str, out_image: str):
     img1_data = img1.get_fdata()
     img1_shape = img1_data.shape
 
-    # If the img2 is a 4D add the forth dimension to the shape of the img1
-    if len(img2.shape) == 4:
-        img1_shape = (img1_shape[0], img1_shape[1], img1_shape[2], img2.shape[3])
-
     # Get data from IM2
     img2_data = img2.get_fdata()
     img2_shape = img2_data.shape
 
+
+    # Multiply the inverse of the affine matrix of img1 by the affine matrix of img2
+    affine_mult = np.linalg.inv(img1_affine) @ img2_affine
+
     # If the img2 is a 4D add the forth dimension to the shape of the img1
     if len(img2_shape) == 4:
         img1_shape = (img1_shape[0], img1_shape[1], img1_shape[2], img2_shape[3])
 
-    # Create an empty array with the same dimensions as IM1
-    new_data = np.zeros(img1_shape)
-
-    for vol in range(img2_data.shape[-1]):
-        # Find the coordinates in voxels of the voxels different from 0 on the img2
-        indices = np.argwhere(img2_data[..., vol] != 0)
+        # Create an empty array with the same dimensions as IM1
+        new_data = np.zeros(img1_shape, dtype=img2_data.dtype)
+
+        for vol in range(img2_data.shape[-1]):
+            # Find the coordinates in voxels of the voxels different from 0 on the img2
+            indices = np.argwhere(img2_data[..., vol] != 0)
+
+            # Apply the affine transformation to the coordinates of the voxels different from 0 on img2
+            new_coords = np.round(affine_mult @ np.concatenate((indices.T, np.ones((1, indices.shape[0]))), axis=0)).astype(int)
 
-        # Multiply the inverse of the affine matrix of img1 by the affine matrix of img2
-        affine_mult = np.linalg.inv(img1_affine) @ img2_affine
+            # Fill the new image with the values of the voxels different from 0 on img2
+            new_data[new_coords[0], new_coords[1], new_coords[2], vol] = img2_data[indices[:, 0], indices[:, 1], indices[:, 2], vol]
+
+    elif len(img2_shape) == 3:
+        # Create an empty array with the same dimensions as IM1
+        new_data = np.zeros(img1_shape, dtype=img2_data.dtype)
+
+        # Find the coordinates in voxels of the voxels different from 0 on the img2
+        indices = np.argwhere(img2_data != 0)
 
         # Apply the affine transformation to the coordinates of the voxels different from 0 on img2
         new_coords = np.round(affine_mult @ np.concatenate((indices.T, np.ones((1, indices.shape[0]))), axis=0)).astype(int)
 
         # Fill the new image with the values of the voxels different from 0 on img2
-        new_data[new_coords[0], new_coords[1], new_coords[2], vol] = img2_data[indices[:, 0], indices[:, 1], indices[:, 2], vol]
+        new_data[new_coords[0], new_coords[1], new_coords[2]] = img2_data[indices[:, 0], indices[:, 1], indices[:, 2]]
 
     # Create a new Nifti image with the same affine and header as IM1
     new_img2 = nib.Nifti1Image(new_data, affine=img1_affine, header=img1.header)

clabtoolkit/segmentationtools.py

@@ -5,13 +5,13 @@
 import numpy as np
 from scipy.spatial import distance
 from scipy.ndimage import convolve
-from scipy.ndimage import label
+from glob import glob
 import nibabel as nib
 from typing import Union
 import clabtoolkit.misctools as cltmisc
 import clabtoolkit.bidstools as cltbids
 import clabtoolkit.imagetools as cltimg
-
+import clabtoolkit.parcellationtools as cltparc
 
 
 def abased_parcellation(t1: str,
@@ -137,28 +137,37 @@ def abased_parcellation(t1: str,
         cmd_cont = cltmisc.generate_container_command(cmd_bashargs, cont_tech, cont_image) # Generating container command
         subprocess.run(cmd_cont, stdout=subprocess.PIPE, universal_newlines=True) # Running container command
 
+        # Deleting the warped images
         cmd_bashargs = ['rm', tmp_xfm_basename + '*Warped.nii.gz']
-        cmd_cont = cltmisc.generate_container_command(cmd_bashargs, cont_tech, cont_image) # Generating container command
-        subprocess.run(cmd_cont, stdout=subprocess.PIPE, universal_newlines=True) # Running container command
+        warped_imgs = glob(tmp_xfm_basename + '*Warped.nii.gz')
+        if len(warped_imgs) > 0:
+            for w_img in warped_imgs:
+                os.remove(w_img)
 
-
+    # Applying spatial transform to the atlas
     if not os.path.isfile(out_parc):
 
         if atlas_type == 'spam':
             # Applying spatial transform
             cmd_bashargs = ['antsApplyTransforms', '-d', '3', '-e', '3', '-i', atlas,
                             '-o', out_parc, '-r', t1, '-t', xfm_invnl,
                             '-t','[' + xfm_affine + ',1]', '-n', interp]
-
+
+            cmd_cont = cltmisc.generate_container_command(cmd_bashargs, cont_tech, cont_image) # Generating container command
+            subprocess.run(cmd_cont, stdout=subprocess.PIPE, universal_newlines=True) # Running container command
+
         elif atlas_type == 'maxprob':
             # Applying spatial transform
-            cmd_bashargs = ['antsApplyTransforms', '-d', '3', '-e', '3', '-i', atlas,
+            cmd_bashargs = ['antsApplyTransforms', '-d', '3', '-i', atlas,
                             '-o', out_parc, '-r', t1, '-t', xfm_invnl,
                             '-t','[' + xfm_affine + ',1]', '-n', 'NearestNeighbor']
 
-        cmd_cont = cltmisc.generate_container_command(cmd_bashargs, cont_tech, cont_image) # Generating container command
-        subprocess.run(cmd_cont, stdout=subprocess.PIPE, universal_newlines=True) # Running container command
-
+            cmd_cont = cltmisc.generate_container_command(cmd_bashargs, cont_tech, cont_image) # Generating container command
+            subprocess.run(cmd_cont, stdout=subprocess.PIPE, universal_newlines=True) # Running container command
+
+            tmp_parc = cltparc.Parcellation(parc_file=out_parc)
+            tmp_parc.save_parcellation(out_file= out_parc, affine=tmp_parc.affine, save_lut=False, save_tsv=False)
+
         # Removing the Warped images
         if os.path.isfile(temp_xfm_invnlw):
             os.remove(temp_xfm_invnlw)