FSL FLIRT Registration with Python & SimpleITK

fsl_registration.py

# Created by Franco Matzkin (vfmatzkin [at] gmail.com)

# This gist gives a basic intuition on how to use tmp files to run FSL's FLIRT
# commands with Python and use those images with the help of SimpleITK.

# Info on how to install FLIRT: https://gist.github.com/vfmatzkin/7987a6edf034d28427696e1fabbb2979 #noqa

import os
import shutil
import tempfile
from multiprocessing import Pool

import SimpleITK as sitk
import nrrd

from .utilities import veri_folder, get_random_string, get_largest_cc
# find them in https://gitlab.com/matzkin/headctools

def find_inverse_mat(mat_file):
    output_file = mat_file.replace(".mat", "_inverse.mat")
    command = (
        "convert_xfm " "-omat {} " "-inverse {}".format(output_file, mat_file)
    )
    os.system(command)
    return output_file


def reverse_registration_folder(
        folder,
        image_ext_inp=".nii.gz",
        image_ext_out=".nrrd",
        interp="nearestneighbour",
        overwrite=False,
):
    commands = []
    for file in os.listdir(folder):
        if file.endswith(image_ext_inp):
            moving_image = os.path.join(folder, file)
            previous_folder, actual_folder = os.path.split(folder)
            orig_folder = os.path.split(previous_folder)[0]
            mat_file = os.path.join(
                previous_folder, file.replace(image_ext_inp, ".mat")
            )
            ref_file = os.path.join(orig_folder, file)
            o_fold = veri_folder(os.path.join(folder, "orig_dims"))
            out_file = os.path.join(o_fold, file)
            if os.path.exists(out_file) and overwrite == False:
                print("{} skipped (already exists).".format(out_file))
                continue
            elif not os.path.exists(mat_file) or not os.path.exists(ref_file):
                print("{} skipped.".format(file))
                continue
            else:
                inverse_mat = find_inverse_mat(mat_file)
                command = apply_mat(
                    moving_image,
                    out_file,
                    inverse_mat,
                    ref_file,
                    interp,
                    return_command=True,
                )
                commands.append(command) if command else 0

    process_pool = Pool(processes=8)
    process_pool.map(execute, commands)


def change_img_ext_folder_sitk(
        ch_path,
        in_ext,
        out_ext,
        out_folder="ext_renamed",
        overwrite=False,
        take_connected_comp=True,
):
    """Change image extensions from a folder using SimpleITK. This function simply opens and saves the files in the
        provided folder with the required extension.

    :param ch_path: Path of the folder containing the images.
    :param in_ext: Extension to be replaced.
    :param out_ext: New extension.
    :return:
    """
    print("[{} -> {}] Folder: {}".format(in_ext, out_ext, ch_path))
    veri_folder(os.path.join(ch_path, out_folder))
    for file in os.listdir(ch_path):
        if file.endswith(in_ext):
            print("    File: {}...".format(file), end="")
            f_path = os.path.join(ch_path, file)
            o_path = os.path.join(ch_path, out_folder, file).replace(
                in_ext, out_ext
            )

            if os.path.exists(o_path) and overwrite == False:
                print(
                    "{} already exists (delete it or activate overwrite flag).".format(
                        o_path
                    )
                )
                continue

            sitk_img = sitk.ReadImage(f_path)

            if take_connected_comp:  # TODO check if always valid
                sitk_img = get_largest_cc(sitk_img)

            sitk.WriteImage(sitk.Cast(sitk_img, sitk.sitkUInt8), o_path)

            nrrd_img = nrrd.read(o_path)
            nrrd_img[1]["encoding"] = "gz"
            nrrd.write(o_path, nrrd_img[0], nrrd_img[1])

            print(" saved in {}.".format(o_path))


def register_fsl(
        moving_image,
        fixed_image,
        out_image=None,
        mat_path=None,
        return_command=False,
        transformation="rigid",
        interp="nearestneighbour",
        overwrite=False,
):
    dof = 6 if transformation == "rigid" else 12  # 9 is affine
    interp = "nearestneighbour" if interp in ['nearestneighbour', 'nearest'] \
        else "trilinear"

    atlas_name = os.path.split(fixed_image)[1].split(".")[0]
    mv_img_p, mv_img_n = os.path.split(moving_image)

    out_dir = os.path.join(mv_img_p, atlas_name)
    veri_folder(out_dir)

    out_image = (
        os.path.join(out_dir, mv_img_n) if out_image is None else out_image
    )

    mat_path = (
        os.path.join(out_dir, mv_img_n.replace(".nii.gz", ".mat"))
        if mat_path is None
        else mat_path
    )

    command = (
        'flirt -in "{}" '
        '-ref "{}" '
        '-out "{}" '
        '-omat "{}" '
        "-bins 256 "
        "-cost corratio "
        "-searchrx -90 90 "
        "-searchry -90 90 "
        "-searchrz -90 90 "
        "-dof {}  "
        "-interp {}".format(
            moving_image, fixed_image, out_image, mat_path, dof, interp
        )
    )
    if return_command:
        return command

    if os.path.exists(out_image) and not overwrite and os.path.getsize(
            out_image) > 0:
        print("\n  The output file already exists (skipping).")
        return

    os.environ["OMP_NUM_THREADS"] = "10"
    os.system(command)


def apply_mat(
        moving_image,
        out_image,
        mat_path,
        reference,
        interp="nearestneighbour",
        return_command=False,
        overwrite=False,
):
    """

    :param moving_image:
    :param out_image:
    :param mat_path:
    :param reference:
    :param interp: nearestneighbour, triliear or sinc
    :param return_command:
    :return:
    """
    if os.path.exists(out_image) and not overwrite:
        print("already exists (skipping).")
        return

    command = (
        "flirt "
        '-in "{}" '
        '-applyxfm -init "{}" '
        '-out "{}" '
        "-paddingsize 0.0 "
        "-interp {} "
        '-ref "{}"'.format(
            moving_image, mat_path, out_image, interp, reference
        )
    )
    if return_command:
        return command

    os.environ["OMP_NUM_THREADS"] = "10"
    os.system(command)
    print(" saved in {}.".format(out_image))
    return 0


def register_folder_flirt(
        folder,
        fixed_image,
        mask_id=None,
        ext=".nii.gz",
        transformation="rigid",
        interp="nearestneighbour",
        overwrite=False,
):
    print("[Registering images to atlas] Folder: {}".format(folder))

    commands_reg = []
    commands_reg_masks = []
    for file in os.listdir(folder):
        if file.endswith(ext) and mask_id not in file:
            print("    File: {}...".format(file), end="")
            moving_image = os.path.join(folder, file)

            command = register_fsl(
                moving_image,
                fixed_image,
                return_command=True,
                transformation=transformation,
                interp=interp,
                overwrite=overwrite,
            )
            reg_img_path = command[
                           command.find("-out") + 6: command.find("-omat") - 2
                           ]
            if overwrite or not os.path.exists(reg_img_path):
                commands_reg.append(command) if command else 0

            if mask_id:
                mat_path = command[
                           command.find("-omat") + 7: command.find("-bins") - 2
                           ]
                mask_file_path = os.path.join(
                    folder, file.replace(ext, "_" + mask_id + ext)
                )

                out_image_name = os.path.split(mask_file_path)[1]
                out_image = os.path.join(
                    os.path.split(mat_path)[0], out_image_name
                )
                if os.path.exists(mask_file_path):
                    comm = apply_mat(
                        mask_file_path,
                        out_image,
                        mat_path,
                        reg_img_path,
                        interp="nearestneighbour",
                        return_command=True,
                        overwrite=overwrite,
                    )
                    if overwrite or not os.path.exists(out_image):
                        commands_reg_masks.append(comm) if comm else 0

    process_pool = Pool(processes=8)
    process_pool.map(execute, commands_reg)
    process_pool.map(execute, commands_reg_masks)


def register_flaps(
        images_folder,
        flaps_folder,
        interp="nearestneighbour",
        out_folder="reg",
        ext=".nii.gz",
):
    print(
        "[Applying .map transforms to folder] Folder: {}".format(flaps_folder)
    )
    veri_folder(os.path.join(flaps_folder, out_folder))

    commands = []
    for file in os.listdir(images_folder):
        if file.endswith(ext):
            print("    File: {}...".format(file), end="")
            moving_image = os.path.join(flaps_folder, file)
            reference = os.path.join(images_folder, file)
            out_image = os.path.join(flaps_folder, out_folder, file)
            mat_file = os.path.join(images_folder, file.replace(ext, ".mat"))

            command = apply_mat(
                moving_image,
                out_image,
                mat_file,
                reference,
                interp,
                return_command=True,
            )
            commands.append(command) if command else 0

    process_pool = Pool(processes=8)
    process_pool.map(execute, commands)


def execute(process):
    os.system(process)


# FSL to SimpleITK

def register_fsl_sitk(moving_image, fixed_image=None, mat_save_path=None,
                      transformation="rigid", interp="nearestneighbour",
                      return_mat=False, mat_to_apply=None, reference=None):
    mv_img = tempfile.NamedTemporaryFile(suffix='.nii.gz', delete=False)
    mv_img_path = mv_img.name

    sitk.WriteImage(moving_image, mv_img_path)

    out_img = tempfile.NamedTemporaryFile(suffix='.nii.gz', delete=False)
    out_img_path = out_img.name

    if fixed_image and type(fixed_image) is not str:
        fx_img = tempfile.NamedTemporaryFile(suffix='.nii.gz', delete=False)
        fx_img_path = fx_img.name
        sitk.WriteImage(fixed_image, fx_img_path)
    elif fixed_image:
        fx_img_path = fixed_image

    if mat_to_apply and reference:
        # Used only when applying .mat transform to mask
        ref_img = tempfile.NamedTemporaryFile(suffix='.nii.gz', delete=False)
        ref_img_path = ref_img.name
        mat = tempfile.NamedTemporaryFile(suffix='.mat', delete=False)
        mat_path = mat.name
        sitk.WriteImage(reference, ref_img_path)
        tmp_mat = open(mat_path, 'w')
        tmp_mat.write(mat_to_apply)
        tmp_mat.close()

        apply_mat(mv_img_path,
                  out_img_path,
                  mat_path,
                  ref_img_path,
                  interp=interp,
                  return_command=False,
                  overwrite=False)
    else:
        register_fsl(
            mv_img_path,
            fx_img_path,
            out_image=out_img_path,
            mat_path=mat_save_path,
            return_command=False,
            transformation=transformation,
            interp=interp
        )

    ret_img = sitk.ReadImage(out_img_path)
    if return_mat:
        mat_f = open(mat_save_path if mat_save_path else mat_to_apply)
        mat_content = mat_f.read()
        mat_f.close()

    if return_mat:
        return ret_img, mat_content
    else:
        return ret_img


if __name__ == "__main__":
    register_folder_flirt(
        "/imgs/all_images/",
        "/imgs/atlas.nii.gz",
        mask_id="mask",
        ext=".nii.gz",
        transformation="rigid",
        interp="nearestneighbour",
    )

    reverse_registration_folder(
        '/imgs/all_images/predictions/')