enh: minor refactor for a more consistent rotation of inputs

Author: oesteban

niworkflows/utils/images.py

@@ -3,14 +3,26 @@
 import numpy as np
 
 
-def as_canonical(img):
-    """Drop rotation w.r.t. cardinal axes of input image."""
+def rotation2canonical(img):
+    """Calculate the rotation w.r.t. cardinal axes of input image."""
     img = nb.as_closest_canonical(img)
-    zooms = list(img.header.get_zooms()[:3])
-    newaff = np.diag(zooms + [1])
-    rot = newaff[:3, :3].dot(np.linalg.inv(img.affine[:3, :3]))
-    newaff[:3, 3] = rot.dot(img.affine[:3, 3])
-    return nb.Nifti1Image(img.dataobj, newaff, img.header)
+    newaff = np.diag(img.header.get_zooms()[:3])
+    r = newaff @ np.linalg.inv(img.affine[:3, :3])
+    if np.allclose(r, np.eye(3)):
+        return None
+    return r
+
+
+def rotate_affine(img, rot=None):
+    """Rewrite the affine of a spatial image."""
+    if rot is None:
+        return img
+
+    img = nb.as_closest_canonical(img)
+    affine = np.eye(4)
+    affine[:3, :3] = rot @ img.affine[:3, :3]
+    affine[:3, 3] = rot @ img.affine[:3, 3]
+    return img.__class__(img.dataobj, affine, img.header)
 
 
 def unsafe_write_nifti_header_and_data(fname, header, data):

niworkflows/viz/utils.py

@@ -20,7 +20,7 @@
 
 from nipype.utils import filemanip
 from .. import NIWORKFLOWS_LOG
-from ..utils.images import as_canonical
+from ..utils.images import rotation2canonical, rotate_affine
 
 
 SVGNS = "http://www.w3.org/2000/svg"
@@ -227,15 +227,17 @@ def plot_segs(
     """
     plot_params = {} if plot_params is None else plot_params
 
-    image_nii = as_canonical(_3d_in_file(image_nii))
-    seg_niis = [as_canonical(_3d_in_file(f)) for f in seg_niis]
+    image_nii = _3d_in_file(image_nii)
+    canonical_r = rotation2canonical(image_nii)
+    image_nii = rotate_affine(image_nii, rot=canonical_r)
+    seg_niis = [rotate_affine(_3d_in_file(f), rot=canonical_r) for f in seg_niis]
     data = image_nii.get_fdata()
 
     plot_params = robust_set_limits(data, plot_params)
 
     bbox_nii = (
         image_nii if bbox_nii is None
-        else as_canonical(_3d_in_file(bbox_nii))
+        else rotate_affine(_3d_in_file(bbox_nii), rot=canonical_r)
     )
 
     if masked: