Convert nifti images to and from RAS+ orientation

RASconverter.py

#!/usr/bin/env python

"""
Function to orient nifti images associated with the KCL BRAIN
centre Bruker scanner, via brukerconvert etc, to and from RAS+
standard orientation.
Requirements:
    - nibabel and numpy libraries
    - a string; name of image to be converted
E MacNicol 20/08/2020 (eilidh.macnicol@kcl.ac.uk)
Update 19/05/2021: include option to append to bids formats and reformatting
"""

# set up #

# import relevant modules
import os
from pathlib import Path
import nibabel as nib
import numpy as np


# set up arguments
def get_parser():
    from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
    """Build parser object."""

    parser = ArgumentParser(
        description='Reorient an image to and from RAS+',
        formatter_class=ArgumentDefaultsHelpFormatter
        )
    parser.add_argument(
        "input",
        help="a string of the input filename",
        action="store",
        type=Path
        )
    parser.add_argument(
        "--output-directory",
        help="path to write output",
        action="store",
        type=Path
        )
    parser.add_argument(
        "--bids",
        help="outputs file in bids format",
        action="store_true",
        default=False
        )
    subparser = parser.add_mutually_exclusive_group()
    subparser.add_argument(
        '--forward',
        help='BRAIN standard orientation -> RAS+ image',
        action="store_true",
        default=False
        )
    subparser.add_argument(
        '--reverse',
        help='RAS+ image -> BRAIN standard orientation',
        action="store_true",
        default=False
        )
    return parser


# functions
def rotate_y(affine):
    print('rotating round y...')
    # set up variables for matrix rotation
    rad = (np.pi)  # 180 degree rotation
    cos_gamma = np.cos(rad)
    sin_gamma = np.sin(rad)

    # define y rotation matrix
    yrot_affine = np.array([
        [cos_gamma,  0,    sin_gamma, 0],
        [0,          1,        0,     0],
        [-sin_gamma, 0,    cos_gamma, 0],
        [0,          0,        0,     1]
    ])
    newaffine = yrot_affine.dot(affine)
    return newaffine


def swap_yz(affine):
    print('swapping axes...')
    newaffine = np.zeros(shape=affine.shape)
    newaffine[0] = affine[0, :]  # keep these axes the same
    newaffine[3] = affine[3, :]  # keep these axes the same
    newaffine[1] = affine[2, :]  # z to y
    newaffine[2] = affine[1, :]  # y to z
    return newaffine


# load image data, header, and affine from args
args = get_parser().parse_args()
img = nib.load(str(args.input.resolve(strict=True)))
data = img.dataobj
hdr = img.header
affine = img.affine

if args.forward:
    new_affine1 = swap_yz(affine)
    new_affine2 = rotate_y(new_affine1)
    newimg = nib.Nifti1Image(data, new_affine2, hdr)
    # change orientation to closest canonical, i.e. RAS+
    finaloutput = nib.as_closest_canonical(newimg)

if args.reverse:
    new_affine1 = rotate_y(affine)
    new_affine2 = swap_yz(new_affine1)
    finaloutput = nib.Nifti1Image(data, new_affine2, hdr)

# save out final image
if args.output_directory is None:
    args.output_directory = os.path.dirname(args.input.resolve())

fileparts = os.path.basename(args.input).split('.')
if args.bids:
    from bids import BIDSLayout
    sfx = BIDSLayout(
        Path(args.input).parents[2]).parse_file_entities(
            filename=os.path.basename(args.input))['suffix']
    if args.forward:
        output = os.path.basename(args.input).replace(
            sfx, str('orientation-RAS_' + sfx))
    elif args.reverse:
        output = os.path.basename(args.input).replace(
            sfx, str('orientation-orig_' + sfx))
else:
    if args.forward:
        if fileparts[0][:-4] != '_RAS':
            output = fileparts[0] + '_RAS.nii.gz'
        else:
            output = fileparts[0]
    elif args.reverse:
        output = fileparts[0] + '_orig.nii.gz'
newname = os.path.join(args.output_directory, output)
nib.nifti1.save(finaloutput, newname)