From ddcbbc8e91377e10a2a38764f6da1436b1e09b96 Mon Sep 17 00:00:00 2001
From: John Wilkie <john@v7labs.com>
Date: Thu, 12 Dec 2024 14:57:27 +0000
Subject: [PATCH] flip axes of legacy imports / exports for `NifTI` dataset
 items

---
 darwin/exporter/formats/nifti.py | 21 +++++++++++++++++++--
 darwin/importer/formats/nifti.py |  6 ++++++
 2 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/darwin/exporter/formats/nifti.py b/darwin/exporter/formats/nifti.py
index 36699b3e5..60b077647 100644
--- a/darwin/exporter/formats/nifti.py
+++ b/darwin/exporter/formats/nifti.py
@@ -164,6 +164,7 @@ def export(
                 image_id=image_id,
                 output_dir=output_dir,
                 legacy=legacy,
+                filename=video_annotation.filename,
             )
 
 
@@ -498,7 +499,7 @@ def unnest_dict_to_list(d: Dict) -> List:
             dataobj=volume.pixel_array.astype(np.int16),
             affine=volume.affine,
         )
-        img = get_reoriented_nifti_image(img, volume)
+        img = get_reoriented_nifti_image(img, volume, legacy, filename)
         if volume.from_raster_layer:
             output_path = Path(output_dir) / f"{image_id}_{volume.class_name}_m.nii.gz"
         else:
@@ -508,16 +509,27 @@ def unnest_dict_to_list(d: Dict) -> List:
         nib.save(img=img, filename=output_path)
 
 
-def get_reoriented_nifti_image(img: nib.Nifti1Image, volume: Dict) -> nib.Nifti1Image:
+def get_reoriented_nifti_image(
+    img: nib.Nifti1Image, volume: Dict, legacy: bool, filename: str
+) -> nib.Nifti1Image:
     """
     Reorients the given NIfTI image based on the original affine.
 
+    For files that require legacy scaling, we flip all axes of the image to be aligned
+    with the target dataset item.
+
     Parameters
     ----------
     img: nib.Nifti1Image
         The NIfTI image to be reoriented
     volume: Dict
         The volume containing the affine and original affine
+    legacy: bool
+        If ``True``, the exporter will flip all axes of the image if the dataset item
+        is not a DICOM
+        If ``False``, the exporter will not flip the axes
+    filename: str
+        The filename of the dataset item
     """
     if volume.original_affine is not None:
         img_ax_codes = nib.orientations.aff2axcodes(volume.affine)
@@ -526,6 +538,11 @@ def get_reoriented_nifti_image(img: nib.Nifti1Image, volume: Dict) -> nib.Nifti1
         orig_ornt = nib.orientations.axcodes2ornt(orig_ax_codes)
         transform = nib.orientations.ornt_transform(img_ornt, orig_ornt)
         img = img.as_reoriented(transform)
+        is_dicom = filename.lower().endswith(".dcm")
+        if legacy and not is_dicom:
+            img = nib.Nifti1Image(
+                np.flip(img.get_fdata(), (0, 1, 2)).astype(np.int16), img.affine
+            )
     return img
 
 
diff --git a/darwin/importer/formats/nifti.py b/darwin/importer/formats/nifti.py
index 888ed36dd..497130ec4 100644
--- a/darwin/importer/formats/nifti.py
+++ b/darwin/importer/formats/nifti.py
@@ -529,6 +529,9 @@ def process_nifti(
     Converts a nifti object of any orientation to the passed ornt orientation.
     The default ornt is LPI.
 
+    For files that require legacy scaling, we flip all axes of the image to be aligned
+    with the target dataset item.
+
     Args:
         input_data: nibabel nifti object.
         ornt: (n,2) orientation array. It defines a transformation to LPI
@@ -547,10 +550,13 @@ def process_nifti(
     orig_ax_codes = nib.orientations.aff2axcodes(img.affine)
     orig_ornt = nib.orientations.axcodes2ornt(orig_ax_codes)
     if remote_file_path in remote_files_that_require_legacy_scaling:
+        is_dicom = remote_file_path.suffix.lower() == ".dcm"
         slot_affine_map = remote_files_that_require_legacy_scaling[remote_file_path]
         affine = slot_affine_map[next(iter(slot_affine_map))]  # Take the 1st slot
         ax_codes = nib.orientations.aff2axcodes(affine)
         ornt = nib.orientations.axcodes2ornt(ax_codes)
+        if not is_dicom:
+            img = nib.Nifti1Image(np.flip(img.get_fdata(), (0, 1, 2)), affine)
     transform = nib.orientations.ornt_transform(orig_ornt, ornt)
     reoriented_img = img.as_reoriented(transform)
     data_array = reoriented_img.get_fdata()