NibabelIOWithReorient saves outputs with wrong orientation in special cases

[mzenk]

During nnunet inference, I noticed that the output orientation of a predicted segmentation is not always identical to the input image orientation.

Here is a demonstration. An example image for which this happens is attached: test_seg.nii.gz

from nnunetv2.imageio.nibabel_reader_writer import NibabelIOWithReorient

test_img_path = "test_seg.nii.gz"
io_fn = NibabelIOWithReorient()
seg, props = io_fn.read_seg(test_img_path)
io_fn.write_seg(seg[0], "test_out.nii.gz", properties=props)

This gives the output:

WARNING: Restored affine does not match original affine. File: test_out.nii.gz
Original affine
 [[ 1.20846860e-01 -9.92671192e-01  0.00000000e+00  1.12361900e+02]
 [ 9.92671192e-01  1.20846860e-01  0.00000000e+00 -1.42911224e+02]
 [ 0.00000000e+00  0.00000000e+00  1.00000000e+00 -1.29000000e+02]
 [ 0.00000000e+00  0.00000000e+00  0.00000000e+00  1.00000000e+00]]
Restored affine
 [[-1.20846860e-01  9.92671192e-01  0.00000000e+00  1.04515481e+02]
 [-9.92671192e-01 -1.20846860e-01  0.00000000e+00 -1.32889562e+02]
 [ 0.00000000e+00  0.00000000e+00  1.00000000e+00 -1.29000000e+02]
 [ 0.00000000e+00  0.00000000e+00  0.00000000e+00  1.00000000e+00]]

This issue is probably related to the discussion here. I will make a PR based on the solution suggested there. As far as I can tell, this only happens in cases where the image axes are rotated wrt to the RAS coordinate system. In such cases, it may be better to use SimpleITKIO, but I think the NibabelIOWithReorient should still handle this special case correctly.

[FabianIsensee]

Thanks a lot!