From ec7a9a4c5b7c007bb60a2358cd0713323144489b Mon Sep 17 00:00:00 2001
From: dorfmanrobert <108150810+dorfmanrobert@users.noreply.github.com>
Date: Fri, 19 Apr 2024 10:55:02 +0100
Subject: [PATCH] [DAR-1649][Internal] Update process_nifti to take custom ornt
 (#818)

* update process_nifti to take custom ornt

* docstring with details
---
 darwin/importer/formats/nifti.py | 33 +++++++++++++++++++-------------
 1 file changed, 20 insertions(+), 13 deletions(-)

diff --git a/darwin/importer/formats/nifti.py b/darwin/importer/formats/nifti.py
index 64e2ee097..c575e029f 100644
--- a/darwin/importer/formats/nifti.py
+++ b/darwin/importer/formats/nifti.py
@@ -470,21 +470,28 @@ def correct_nifti_header_if_necessary(img_nii):
 
 
 def process_nifti(
-    input_data: Union[Sequence[nib.nifti1.Nifti1Image], nib.nifti1.Nifti1Image]
-):
+    input_data: nib.nifti1.Nifti1Image,
+    ornt: Optional[List[List[float]]] = [[0.0, -1.0], [1.0, -1.0], [2.0, -1.0]],
+) -> Tuple[np.ndarray, Tuple[float]]:
     """
-    Function which takes in a single nifti path or a list of nifti paths
-    and returns the pixel_array, affine and pixdim
+    Function that converts a nifti object to RAS orientation, then converts to the passed ornt orientation.
+    The default ornt is for LPI.
+
+    Args:
+        input_data: nibabel nifti object.
+        ornt: (n,2) orientation array.
+            ornt[N,1] is a flip of axis N of the array, where 1 means no flip and -1 means flip.
+            ornt[:,0] is the transpose that needs to be done to the implied array, as in arr.transpose(ornt[:,0]).
+
+    Returns:
+        data_array: pixel array with orientation ornt.
+        pixdims: tuple of nifti header zoom values.
     """
-    if isinstance(input_data, nib.nifti1.Nifti1Image):
-        img = correct_nifti_header_if_necessary(input_data)
-        img = nib.funcs.as_closest_canonical(img)
-        nib.orientations.aff2axcodes(img.affine)
-        # TODO: Future feature to pass custom ornt could go here.
-        ornt = [[0.0, -1.0], [1.0, -1.0], [1.0, -1.0]]
-        data_array = nib.orientations.apply_orientation(img.get_fdata(), ornt)
-        pixdims = img.header.get_zooms()
-        return data_array, pixdims
+    img = correct_nifti_header_if_necessary(input_data)
+    img = nib.funcs.as_closest_canonical(img)
+    data_array = nib.orientations.apply_orientation(img.get_fdata(), ornt)
+    pixdims = img.header.get_zooms()
+    return data_array, pixdims
 
 
 def convert_to_dense_rle(raster: np.ndarray) -> List[int]: