Medical-Tools/reorient_to_rai.py at main · bigorgs/Medical-Tools · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
import argparse
from glob import glob
import os
import numpy as np
import itk


def reorient_to_rai(image):
    """
    Reorient image to RAI orientation.
    :param image: Input itk image.
    :return: Input image reoriented to RAI.
    """
    filter = itk.OrientImageFilter.New(image)
    filter.UseImageDirectionOn()
    filter.SetInput(image)
    m = itk.GetMatrixFromArray(np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], np.float64))
    filter.SetDesiredCoordinateDirection(m)
    filter.Update()
    reoriented = filter.GetOutput()
    return reoriented


if __name__ == '__main__':
    # parser = argparse.ArgumentParser()
    # parser.add_argument('--image_folder', default="./output" , type=str, required=True)
    # parser.add_argument('--output_folder', default="./output" ,type=str, required=True)
    # parser_args = parser.parse_args()
    # if not os.path.exists(parser_args.output_folder):
    #     os.makedirs(parser_args.output_folder)

    image_folder="./output"
    output_folder="./output"
    # filenames = glob(os.path.join(parser_args.image_folder, '*.nii.gz'))
    filenames = glob(os.path.join(image_folder, '*.nii.gz'))
    for filename in sorted(filenames):
        basename = os.path.basename(filename)
        basename_wo_ext = basename[:basename.find('.nii.gz')]
        print(basename_wo_ext)
        image = itk.imread(filename)
        reoriented = reorient_to_rai(image)

        reoriented.SetOrigin([0, 0, 0])
        m = itk.GetMatrixFromArray(np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], np.float64))
        reoriented.SetDirection(m)
        reoriented.Update()
        # itk.imwrite(reoriented, os.path.join(parser_args.output_folder, basename_wo_ext + '1.nii.gz'))
        itk.imwrite(reoriented, os.path.join(output_folder, basename_wo_ext + '2.nii.gz'))