Updated imageOP with warp related functions

src/imageOp.py

@@ -20,7 +20,7 @@
 import nibabel
 import pandas as pd
 from nilearn.input_data import NiftiMasker
-
+import fileOP as fop
 
 def get_dimensions(nifti_file: str) -> int:
     """
@@ -89,3 +89,126 @@ def mask_img(nifti_file: str, mask_file: str, masked_file: str) -> str:
          range(img.GetNumberOfComponentsPerPixel())])
     SimpleITK.WriteImage(SimpleITK.Cast(masked_img, SimpleITK.sitkVectorFloat32), masked_file)
     return masked_file
+
+
+def get_average_warps(warp_dir: str, w_card: str, out_dir: str = None) -> sitk.Image:
+    """
+    Computes the average of multiple warp images. The images must be multi-component ones.
+
+    :rtype: sitk.Image
+    :param warp_dir: 
+    :param w_card: 
+    :param out_dir: 
+    :return: The average multi-component warp image.
+    """
+    image_paths = fop.get_files(warp_dir, w_card)
+
+    print(f"Starting with {image_paths[0]}")
+    summed_image = sitk.ReadImage(image_paths[0], sitk.sitkVectorFloat64)
+
+    # Sum all other images
+    for image_index in range(1, len(image_paths)):
+        print(f"       Adding {image_paths[image_index]} to the image stack.")
+        current_image = sitk.ReadImage(image_paths[image_index], sitk.sitkVectorFloat64)
+        summed_image = summed_image + current_image
+
+    # Compute the mean
+    number_of_images = len(image_paths)
+    avg_val = 1 / number_of_images
+    average_image = sitk.Compose([sitk.VectorIndexSelectionCast(summed_image, i) * avg_val for i in
+                                  range(summed_image.GetNumberOfComponentsPerPixel())])
+
+    if out_dir is not None:
+        sitk.WriteImage(average_image, os.path.join(out_dir, "average_warp.nii.gz"))
+    return average_image
+
+
+def get_magnitude_warp(input_warp_path: str, magnitude_warp_path: str = None) -> sitk.Image:
+    """
+    Computes the voxel wise magnitude of a multi-component warp image.
+
+    :rtype: sitk.Image
+    :param input_warp_path: The path to the warp image
+    :param magnitude_warp_path: The (optional) desired path where the magnitude image will be stored
+    :return: The magnitude image.
+    """
+    # Read multi component image
+    image = sitk.ReadImage(input_warp_path, sitk.sitkVectorFloat64)
+    # Compute magnitude image
+    magnitude_image = sitk.VectorMagnitude(image)
+    # Write to location is specified
+    if magnitude_warp_path is not None:
+        sitk.WriteImage(magnitude_image, magnitude_warp_path)
+    return magnitude_image
+
+
+def get_mean_std_magnitude(magnitude_dir: str, w_card: str, out_dir: str = None) -> (sitk.Image, sitk.Image):
+    """
+    Computes the mean and std of multiple magnitude images. The images must be single-component ones.
+
+    :rtype: (sitk.Image, sitk.Image)
+    :param magnitude_dir: The directory where all magnitude images are located
+    :param w_card: The pattern to retrieve the magnitude images from the folder at magnitude_dir
+    :param out_dir: The (optional) desired path where the mean and std image will be stored.
+    :return: The mean and the std image.
+    """
+    image_paths = fop.get_files(magnitude_dir, w_card)
+
+    # Compute the mean
+    print(f"Computing the Mean")
+    print(f"Starting with {image_paths[0]}")
+    summed_image = sitk.ReadImage(image_paths[0])
+
+    # Sum all other images
+    for image_index in range(1, len(image_paths)):
+        print(f"       Adding {image_paths[image_index]} to the image stack.")
+        current_image = sitk.ReadImage(image_paths[image_index])
+        summed_image = summed_image + current_image
+
+    number_of_images = len(image_paths)
+    avg_val = 1 / number_of_images
+    mean_image = summed_image * avg_val
+
+    # Compute the STD
+    print(f"Computing the STD.")
+    print(f"Starting with {image_paths[0]}")
+    diff_image = sitk.ReadImage(image_paths[0]) - mean_image
+    summed_image = diff_image * diff_image
+
+    # Process all other images
+    for image_index in range(1, len(image_paths)):
+        print(f"       Adding {image_paths[image_index]} to the image stack.")
+        diff_image = sitk.ReadImage(image_paths[0]) - mean_image
+        summed_image = summed_image + diff_image
+    std_image = pow(summed_image * 1 / number_of_images, 1/2)
+
+    if out_dir is not None:
+        print(f"Writing mean and std image at {out_dir}.")
+        sitk.WriteImage(mean_image, os.path.join(out_dir, "mean_warp.nii.gz"))
+        sitk.WriteImage(std_image, os.path.join(out_dir, "std_warp.nii.gz"))
+
+    return mean_image, std_image
+
+
+def get_mean_std_magnitude_warps(input_directory: str, wildcard: str = "*warp.nii.gz") -> (sitk.Image, sitk.Image):
+    """
+    Computes a magnitude image for each file that is matched by wildcard in input_directory together with the mean 
+    and std of all images. 
+
+    :rtype: (sitk.Image, sitk.Image)
+    :param input_directory: The path to the directory where all warps are located
+    :param wildcard: The pattern to retrieve the magnitude images from the folder at input_directory
+    :return: The mean and std image.
+    """
+    # Get files
+    files = fop.get_files(input_directory, wildcard)
+
+    # Get magnitude for each file
+    for file in files:
+        magnitude_file_name = os.path.basename(file)
+        print(f"Computing magnitude image ({magnitude_file_name}) for {file}")
+        get_magnitude_warp(file, os.path.join(input_directory, f"{magnitude_file_name}_magnitude.nii.gz"))
+
+    # Get mean and std magnitude images
+    mean_magnitude, std_magnitude = get_mean_std_magnitude(input_directory, "*_magnitude.nii.gz", input_directory)
+    return mean_magnitude, std_magnitude