Fix MRDS output corruption and orientation handling

src/scilpy/cli/scil_mrds_metrics.py

@@ -10,14 +10,14 @@
      FA of tensor D_3 will be in index 2 of the 4th dimension.
 """
 
-import logging
-import numpy as np
-import nibabel as nib
 import argparse
+import logging
 
+import numpy as np
 from dipy.reconst.dti import fractional_anisotropy
 
 from scilpy.io.image import get_data_as_mask
+from scilpy.io.stateful_image import StatefulImage
 from scilpy.io.utils import (add_overwrite_arg,
                              add_verbose_arg,
                              assert_inputs_exist, assert_outputs_exist,
@@ -80,17 +80,16 @@ def main():
     assert_outputs_exist(parser, args, [],
                          optional=[args.fa, args.ad, args.rd, args.md])
 
-    evals_img = nib.load(args.in_evals)
-    lambdas = evals_img.get_fdata(dtype=np.float32)
-
-    header = evals_img.header
-    affine = evals_img.affine
+    evals_simg = StatefulImage.load(args.in_evals)
+    lambdas = evals_simg.get_fdata(dtype=np.float32)
 
     X, Y, Z = lambdas.shape[0:3]
 
     # load mask
     if args.mask:
-        mask = get_data_as_mask(nib.load(args.mask))
+        mask_simg = StatefulImage.load(args.mask)
+        mask_simg.reorient(evals_simg.axcodes)
+        mask = get_data_as_mask(mask_simg)
     else:
         mask = np.ones((X, Y, Z), dtype=np.uint8)
 
@@ -104,44 +103,36 @@ def main():
                        fractional_anisotropy(lambdas[:, :, :, 3:6]),
                        fractional_anisotropy(lambdas[:, :, :, 6:9])),
                       axis=3)
-        nib.save(nib.Nifti1Image(fa * mask[..., None],
-                                 affine=affine,
-                                 header=header,
-                                 dtype=np.float32),
-                 args.fa)
+        StatefulImage.create_from(
+            (fa * mask[..., None]).astype(np.float32),
+            evals_simg, is_orientation=False).save(args.fa)
 
     if args.ad:
         ad = np.stack((lambdas[:, :, :, 0],
                        lambdas[:, :, :, 3],
                        lambdas[:, :, :, 6]),
                       axis=3)
-        nib.save(nib.Nifti1Image(ad * mask[..., None],
-                                 affine=affine,
-                                 header=header,
-                                 dtype=np.float32),
-                 args.ad)
+        StatefulImage.create_from(
+            (ad * mask[..., None]).astype(np.float32),
+            evals_simg, is_orientation=False).save(args.ad)
 
     if args.rd:
         rd = np.stack(((lambdas[:, :, :, 1] + lambdas[:, :, :, 2])/2,
                        (lambdas[:, :, :, 4] + lambdas[:, :, :, 5])/2,
                        (lambdas[:, :, :, 7] + lambdas[:, :, :, 8])/2),
                       axis=3)
-        nib.save(nib.Nifti1Image(rd * mask[..., None],
-                                 affine=affine,
-                                 header=header,
-                                 dtype=np.float32),
-                 args.rd)
+        StatefulImage.create_from(
+            (rd * mask[..., None]).astype(np.float32),
+            evals_simg, is_orientation=False).save(args.rd)
 
     if args.md:
         md = np.stack((np.average(lambdas[:, :, :, 0:3], axis=3),
                        np.average(lambdas[:, :, :, 3:6], axis=3),
                        np.average(lambdas[:, :, :, 6:9], axis=3)),
                       axis=3)
-        nib.save(nib.Nifti1Image(md * mask[..., None],
-                                 affine=affine,
-                                 header=header,
-                                 dtype=np.float32),
-                 args.md)
+        StatefulImage.create_from(
+            (md * mask[..., None]).astype(np.float32),
+            evals_simg, is_orientation=False).save(args.md)
 
 
 if __name__ == '__main__':

src/scilpy/cli/scil_mrds_select_number_of_tensors.py

@@ -31,11 +31,11 @@
 import itertools
 import logging
 
-import nibabel as nib
 import numpy as np
 
 from scilpy.image.labels import get_data_as_labels
 from scilpy.io.image import get_data_as_mask
+from scilpy.io.stateful_image import StatefulImage
 from scilpy.io.utils import (add_overwrite_arg, add_processes_arg,
                              add_sh_basis_args, add_verbose_arg,
                              assert_headers_compatible,
@@ -89,29 +89,43 @@ def main():
     assert_outputs_exist(parser, args, output_files)
     assert_headers_compatible(parser, [args.in_volume] + [x for xs in mrds_files for x in xs])
 
+    # MOdel SElector MAP
+    mosemap_simg = StatefulImage.load(args.in_volume)
+    mosemap = get_data_as_labels(mosemap_simg)
+    X, Y, Z = mosemap.shape[0:3]
+
     signal_fraction = []
     evals = []
     iso = []
     num_tensors = []
     evecs = []
     for N in range(3):
-        signal_fraction.append(nib.load(mrds_files[N][0]).get_fdata(dtype=np.float32))
-        evals.append(nib.load(mrds_files[N][1]).get_fdata(dtype=np.float32))
-        iso.append(nib.load(mrds_files[N][2]).get_fdata(dtype=np.float32))
-        num_tensors.append(nib.load(mrds_files[N][3]).get_fdata(dtype=np.float32))
-        evecs.append(nib.load(mrds_files[N][4]).get_fdata(dtype=np.float32))
+        sf_simg = StatefulImage.load(mrds_files[N][0])
+        sf_simg.reorient(mosemap_simg.axcodes)
+        signal_fraction.append(sf_simg.get_fdata(dtype=np.float32))
 
-    # MOdel SElector MAP
-    mosemap_img = nib.load(args.in_volume)
-    mosemap = get_data_as_labels(mosemap_img)
-    header = mosemap_img.header
+        evals_simg = StatefulImage.load(mrds_files[N][1])
+        evals_simg.reorient(mosemap_simg.axcodes)
+        evals.append(evals_simg.get_fdata(dtype=np.float32))
 
-    affine = mosemap_img.affine
-    X, Y, Z = mosemap.shape[0:3]
+        iso_simg = StatefulImage.load(mrds_files[N][2])
+        iso_simg.reorient(mosemap_simg.axcodes)
+        iso.append(iso_simg.get_fdata(dtype=np.float32))
+
+        num_tensors_simg = StatefulImage.load(mrds_files[N][3])
+        num_tensors_simg.reorient(mosemap_simg.axcodes)
+        num_tensors.append(num_tensors_simg.get_fdata(dtype=np.float32))
+
+        # evecs are stored in world space. A plain grid reslice does not change their values.
+        evecs_simg = StatefulImage.load(mrds_files[N][4])
+        evecs_simg.reorient(mosemap_simg.axcodes)
+        evecs.append(evecs_simg.get_fdata(dtype=np.float32))
 
     # load mask
     if args.mask:
-        mask = get_data_as_mask(nib.load(args.mask), dtype=bool)
+        mask_simg = StatefulImage.load(args.mask)
+        mask_simg.reorient(mosemap_simg.axcodes)
+        mask = get_data_as_mask(mask_simg, dtype=bool)
     else:
         mask = np.ones((X, Y, Z), dtype=np.uint8)
 
@@ -141,26 +155,22 @@ def main():
             evecs_out[X, Y, Z, :] = evecs[N][X, Y, Z, :]
 
     # write output files
-    nib.save(nib.Nifti1Image(signal_fraction_out,
-                             affine=affine,
-                             header=header,
-                             dtype=np.float32), output_files[0])
-    nib.save(nib.Nifti1Image(evals_out,
-                             affine=affine,
-                             header=header,
-                             dtype=np.float32), output_files[1])
-    nib.save(nib.Nifti1Image(iso_out,
-                             affine=affine,
-                             header=header,
-                             dtype=np.float32), output_files[2])
-    nib.save(nib.Nifti1Image(num_tensors_out,
-                             affine=affine,
-                             header=header,
-                             dtype=np.uint8), output_files[3])
-    nib.save(nib.Nifti1Image(evecs_out,
-                             affine=affine,
-                             header=header,
-                             dtype=np.float32), output_files[4])
+    # evecs are in world space. No direction rotation is required on save.
+    StatefulImage.create_from(
+        signal_fraction_out.astype(np.float32), mosemap_simg,
+        is_orientation=False).save(output_files[0])
+    StatefulImage.create_from(
+        evals_out.astype(np.float32), mosemap_simg,
+        is_orientation=False).save(output_files[1])
+    StatefulImage.create_from(
+        iso_out.astype(np.float32), mosemap_simg,
+        is_orientation=False).save(output_files[2])
+    StatefulImage.create_from(
+        num_tensors_out.astype(np.uint8), mosemap_simg,
+        is_orientation=False).save(output_files[3])
+    StatefulImage.create_from(
+        evecs_out.astype(np.float32), mosemap_simg,
+        is_orientation=False).save(output_files[4])
 
 
 if __name__ == '__main__':

src/scilpy/cli/tests/test_mrds_select_number_of_tensors.py

@@ -4,8 +4,13 @@
 import os
 import tempfile
 
+import nibabel as nib
+import numpy as np
+from PIL import Image
+
 from scilpy import SCILPY_HOME
 from scilpy.io.fetcher import fetch_data, get_testing_files_dict
+from scilpy.io.stateful_image import StatefulImage
 
 # If they already exist, this only takes 5 seconds (check md5sum)
 fetch_data(get_testing_files_dict(), keys=['mrds.zip'])
@@ -44,3 +49,88 @@ def test_execution_mrds_w_mask(script_runner, monkeypatch):
                              '--mask', in_mask,
                              '-f'])
     assert ret.success
+
+
+def test_non_ras_mrds_select_number_of_tensors(script_runner, monkeypatch):
+    monkeypatch.chdir(os.path.expanduser(tmp_dir.name))
+
+    nufo_path = os.path.join(SCILPY_HOME, 'mrds', 'sub-01_nufo.nii.gz')
+    nufo_img = nib.load(nufo_path)
+    prefix = os.path.join(SCILPY_HOME, 'mrds', 'sub-01')
+
+    # sub-01 MRDS data is in LAS orientation.
+    # Create an axial mid-slice mask in LAS.
+    mid_z = nufo_img.shape[2] // 2
+    mask_las = np.zeros(nufo_img.shape, dtype=np.uint8)
+    mask_las[:, :, mid_z] = 1
+    las_mask_path = 'mask_midslice_las.nii.gz'
+    nib.save(nib.Nifti1Image(mask_las, nufo_img.affine), las_mask_path)
+
+    # Reorient mask to coronal/sagittal axis order (PIR)
+    simg_mask = StatefulImage.convert_to_simg(
+        nib.Nifti1Image(mask_las, nufo_img.affine))
+    simg_mask.reorient(('P', 'I', 'R'))
+    pir_mask_path = 'mask_midslice_pir.nii.gz'
+    nib.save(nib.Nifti1Image(simg_mask.get_fdata().astype(np.uint8),
+                             simg_mask.affine),
+             pir_mask_path)
+
+    # Reference run with native LAS mask
+    ret_ref = script_runner.run([
+        'scil_mrds_select_number_of_tensors', prefix, nufo_path,
+        '--mask', las_mask_path,
+        '--out_prefix', 'out_ref', '-f'])
+    assert ret_ref.success
+
+    # Test run with cross-orientation PIR mask
+    ret_test = script_runner.run([
+        'scil_mrds_select_number_of_tensors', prefix, nufo_path,
+        '--mask', pir_mask_path,
+        '--out_prefix', 'out_test', '-f'])
+    assert ret_test.success
+
+    # Verify LAS output orientation is preserved
+    out_num_img = nib.load('out_test_MRDS_num_tensors.nii.gz')
+    assert nib.orientations.aff2axcodes(out_num_img.affine) == ('L', 'A', 'S')
+    out_num = out_num_img.get_fdata()
+
+    # In LAS, the mid-slice along z should be kept (non-zero),
+    # while other slices should be masked out (zero).
+    assert np.sum(out_num[:, :, mid_z]) > 0
+    assert np.all(out_num[:, :, mid_z - 1] == 0)
+
+    # Verify float output is not quantized to uint8 despite
+    # uint8 NUFO reference
+    out_sf = nib.load('out_test_MRDS_signal_fraction.nii.gz')
+    assert out_sf.header.get_data_dtype() == np.float32
+    sf_data = out_sf.get_fdata()
+    assert np.any((sf_data > 0) & (sf_data < 1))
+
+    # Verify rendered peaks screenshots match reference
+    ret_viz_ref = script_runner.run([
+        'scil_viz_volume_screenshot', nufo_path, 'ref_peaks.png',
+        '--peaks', 'out_ref_MRDS_evecs.nii.gz',
+        '--slices', str(mid_z),
+        '--volume_opacity', '0.0', '-f'])
+    assert ret_viz_ref.success
+
+    ret_viz_test = script_runner.run([
+        'scil_viz_volume_screenshot', nufo_path, 'test_peaks.png',
+        '--peaks', 'out_test_MRDS_evecs.nii.gz',
+        '--slices', str(mid_z),
+        '--volume_opacity', '0.0', '-f'])
+    assert ret_viz_test.success
+
+    ref_arr = np.asarray(
+        Image.open(f'ref_peaks_slice_{mid_z}.png')).astype(np.float32)
+    test_arr = np.asarray(
+        Image.open(f'test_peaks_slice_{mid_z}.png')).astype(np.float32)
+
+    # Confirm peak glyphs are rendered with non-trivial pixel variance
+    assert np.std(test_arr) > 3.0
+
+    # Confirm visual equivalence between reference and reoriented mask
+    diff = np.abs(ref_arr - test_arr)
+    assert np.allclose(ref_arr, test_arr, atol=1.0), (
+        f"Peaks render with reoriented mask differs from reference "
+        f"(max diff = {np.max(diff):.2f})")