From 20d0a306341b461e15aa8c6e782a3cfb86cdf78a Mon Sep 17 00:00:00 2001
From: John Wilkie <john@v7labs.com>
Date: Sun, 3 Nov 2024 15:39:17 +0000
Subject: [PATCH] Fixed NifTI import orientation

---
 darwin/importer/formats/nifti.py              | 15 +++++++--
 .../importer/formats/import_nifti_test.py     | 31 ++++++++++++++++++-
 2 files changed, 43 insertions(+), 3 deletions(-)

diff --git a/darwin/importer/formats/nifti.py b/darwin/importer/formats/nifti.py
index a86a70d89..8f4722164 100644
--- a/darwin/importer/formats/nifti.py
+++ b/darwin/importer/formats/nifti.py
@@ -517,7 +517,7 @@ def process_nifti(
 ) -> Tuple[np.ndarray, Tuple[float]]:
     """
     Function that converts a nifti object to RAS orientation (if legacy), then converts to the passed ornt orientation.
-    The default ornt is for LPI.
+    The default ornt is LPI.
 
     Args:
         input_data: nibabel nifti object.
@@ -532,7 +532,18 @@ def process_nifti(
     img = correct_nifti_header_if_necessary(input_data)
     if legacy:
         img = nib.funcs.as_closest_canonical(img)
-    data_array = nib.orientations.apply_orientation(img.get_fdata(), ornt)
+        data_array = nib.orientations.apply_orientation(img.get_fdata(), ornt)
+    else:
+        current_ornt = nib.orientations.io_orientation(img.affine)
+        adjusted_ornt = []
+        for idx, (target_axis, target_direction) in enumerate(ornt):
+            _, current_direction = current_ornt[idx]
+            if current_direction == target_direction:
+                adjusted_ornt.append([target_axis, 1])
+            else:
+                adjusted_ornt.append([target_axis, -1])
+        data_array = nib.orientations.apply_orientation(img.get_fdata(), adjusted_ornt)
+
     pixdims = img.header.get_zooms()
     return data_array, pixdims
 
diff --git a/tests/darwin/importer/formats/import_nifti_test.py b/tests/darwin/importer/formats/import_nifti_test.py
index d2c397e48..58cd6a28b 100644
--- a/tests/darwin/importer/formats/import_nifti_test.py
+++ b/tests/darwin/importer/formats/import_nifti_test.py
@@ -9,6 +9,7 @@
 import numpy as np
 import pytest
 from scipy import ndimage
+import nibabel as nib
 
 from darwin.datatypes import (
     Annotation,
@@ -17,7 +18,7 @@
     SubAnnotation,
     VideoAnnotation,
 )
-from darwin.importer.formats.nifti import get_new_axial_size, parse_path
+from darwin.importer.formats.nifti import get_new_axial_size, parse_path, process_nifti
 from tests.fixtures import *
 
 
@@ -238,6 +239,34 @@ def test_get_new_axial_size_with_isotropic():
     assert new_size == (20, 10)
 
 
+def test_process_nifti_orientation_legacy(team_slug_darwin_json_v2):
+    pass
+    # It's not yet clear how to write this as the desired behaviour isn't known
+
+
+def test_process_nifti_orinetation_no_legacy(team_slug_darwin_json_v2):
+    with tempfile.TemporaryDirectory() as tmpdir:
+        with ZipFile("tests/data.zip") as zfile:
+            zfile.extractall(tmpdir)
+            filepath = (
+                Path(tmpdir)
+                / team_slug_darwin_json_v2
+                / "nifti"
+                / "releases"
+                / "latest"
+                / "annotations"
+                / "vol0_brain.nii.gz"
+            )
+            lpi_ornt = [[0.0, -1.0], [1.0, -1.0], [2.0, -1.0]]
+            ras_file = nib.load(filepath)
+            las_transformed_file = nib.orientations.apply_orientation(
+                ras_file.get_fdata(), lpi_ornt
+            )
+            processed_file, _ = process_nifti(input_data=ras_file, legacy=False)
+            assert not np.array_equal(processed_file, ras_file._dataobj)
+            assert np.array_equal(processed_file, las_transformed_file)
+
+
 def serialise_annotation_file(
     annotation_file: AnnotationFile, as_dict
 ) -> Union[str, dict]: