B6_dapi_template/template_creation_pipeline/template_creation.py

from helper_functions import files_in_dir, make_dirs, add_slash_to_path, flip_volume_lr, robust_avg_volumes,\
    files_ind_dir_wo_string, avg_template
import sys
import getopt
import os
import shutil
from nipype.interfaces import ants


# Help message.
tab = '     '
usage_message = '\nCOMMAND:\n' + tab + sys.argv[0] + \
                '\n\nOPTIONS:\n' + \
                tab + '-h, help menu\n' + \
                tab + '-i, --images, <list of input images> or path to directory with only .nii images.\n' + \
                tab + '-o, --output, <path to output directory>\n' + \
                tab + '-s, --symmetric, [bool] use symmetric modelling. Default: False. Currently only works for 3d images.\n' + \
                tab + '-t, --template, <path to template> can include template. Will not be included in final average. \n' + \
                tab + '-w, --weight, [Double] specifies the weight a given template will have in averages. From 0 to 1. Default: 1/(number of images + 1)\n' + \
                tab + '-l, --log, [bool]>\n'


# Registration
def registration(fixed, moving, result, reg_type, sigma, shrink, param, dim, initial_trans='', inverse_result=''):
    """
    Registration function. Registeres the moving image to the fixed with given parameters. If initial transform is not
    given then the images are aligned by center of mass. Outputs inversely transformed result if given.
    :param fixed: Path to fixed image.
    :param moving: Path to moving image.
    :param result: Path to result image.
    :param reg_type: Registration type. List containing, Rigid, Affine, and/or SyN.
    :param sigma: List of sigma parameters. Ordered according to the reg_type parameter.
    :param shrink: List of shrink parameters. Ordered according to the reg_type parameter.
    :param param: List of transform parameters. Ordered according to the reg_type parameter.
    :param dim: Dimension of images. Typically 2 or 3.
    :param initial_trans: Optional path to initial moving transform. If not given, the images will be matched initially
    by aligning the center of mass.
    :param inverse_result: Optional path to the inversely transformed fixed image.
    :return: Returns nothing.
    """
    # Extract number of registrations.
    steps = len(reg_type)

    # Add 1000 iterations for each step.
    if steps == 1:
        iterations = [[1000] * len(sigma[0])]
    else:
        iterations = []
        for i, reg in enumerate(reg_type):
            iteration = [1000] * len(sigma[i])
            iterations.append(iteration)

    # Create registration instance.
    reg_instance = ants.Registration(dimension=dim,
                                     transforms=reg_type,
                                     transform_parameters=param,
                                     metric=['MI']*steps,
                                     interpolation='BSpline',
                                     fixed_image=[fixed],
                                     moving_image=[moving],
                                     metric_weight=[1]*steps,
                                     radius_or_number_of_bins=[32]*steps,
                                     number_of_iterations=iterations,
                                     smoothing_sigmas=sigma,
                                     shrink_factors=shrink,
                                     convergence_threshold=[1.e-7],
                                     sampling_strategy=['Regular']*steps,
                                     use_histogram_matching=True,
                                     winsorize_lower_quantile=0.05,
                                     winsorize_upper_quantile=0.95,
                                     output_warped_image=result,
                                     write_composite_transform=True,
                                     output_transform_prefix=result.split('.')[0] + '-Transformation',
                                     num_threads=12)
    # Add initial moving transform if given, else match by center of mass.
    if not initial_trans == '':
        reg_instance.inputs.initial_moving_transform = initial_trans
    else:
        reg_instance.inputs.initial_moving_transform_com = 1

    # Output reverse results if path is given.
    if not inverse_result == '':
        reg_instance.inputs.output_inverse_warped_image = inverse_result

    # Run registration.
    reg_instance.run()


def create_template_from_images(images_input, results_dir, symmetric=False, template='', template_weight=1,
                                r=3, a=2, nl=3, print_log=False):
    """
    Population based template creation from input images. If input template is given this will be part of the averaging
    with weight template_weight. Will perform r rigid, a rigid and affine, and nl rigid, affine, and non-linear
    iterations. Will flip all input images in the left-right plane if symmetric is set to true.
    :param images_input: List of paths to images.
    :param results_dir: Working directory.
    :param symmetric: Boolean [Default: False]. Will flip all input images in the left-right plane if symmetric is set
    to true.
    :param template: str [Default: '']. Will include template in averaging if path is specified.
    :param template_weight: float [Default: 1]. Will weight the template with given weight if set to different from 1.
    Can be in the range of 0 < template_weight < 1.
    :param r: Int [Default: 3]. Number of rigid registrations.
    :param a: Int [Default: 2]. Number of rigid and affine registrations.
    :param nl: Int [Default: 3]. Number of rigid, affine, and non-linear registrations.
    :param print_log: Boolean [Default: False]. Will print log if true.
    :return: Returns nothing.
    """
    out_dir = add_slash_to_path(results_dir)
    make_dirs(out_dir)
    template_present = not template == ''

    if print_log:
        print('Determining images-variable format and defining input images.')
    if isinstance(images_input, list):
        images = images_input
    elif isinstance(images_input, str):
        images_input = add_slash_to_path(images_input)
        images = files_in_dir(images_input, '.nii')

    all_images = images

    if symmetric:
        flipped_path = out_dir + 'raw_flipped/'
        flipped_images = []
        make_dirs(flipped_path)
        if print_log:
            print('Creating flipped images.')
        for image in images:
            flipped_image = flipped_path + image.split('/')[-1].split('.')[0] + '_flipped.nii'
            if not os.path.exists(flipped_image):
                flip_volume_lr(image, flipped_image)
            else:
                if print_log:
                    print('Flipped image', flipped_image, 'already present.')
            flipped_images.append(flipped_image)

        all_images.extend(flipped_images)

    if print_log:
        print('Defining iterations.')
    reg_types = ['Rigid'] * r + ['Affine'] * a + ['SyN'] * nl

    if print_log:
        print(reg_types)

    for i, iteration in enumerate(reg_types):
        iteration_nr = str(i).zfill(2)
        cur_it_dir = out_dir + iteration_nr + '/'
        make_dirs(cur_it_dir)

        if i == 0:
            robust_avg_volumes(all_images, out_dir + '00-average.nii')
            if template_present:
                avg_template(out_dir + '00-average.nii',
                             template,
                             template_weight,
                             out_dir + '00-average-template.nii')
        else:
            shutil.rmtree(out_dir + str(i-1).zfill(2))

        average = out_dir + iteration_nr + '-average-template.nii' if template_present else out_dir + iteration_nr + '-average.nii'

        # Iteration specific registration parameters.
        if iteration == 'Rigid':
            reg = ['Rigid']
            params = [(0.25,)]
            sigma = [[4, 4, 2, 2]]
            shrink = [[32, 16, 8, 4]]
            if i == 0:
                sigma = [[4, 4, 4]]
                shrink = [[32, 16, 8]]
        elif iteration == 'Affine':
            reg = ['Rigid', 'Affine']
            params = [(0.25,), (0.25,), ]
            sigma = [[4, 4, 2, 2], [4, 4, 2, 2]]
            shrink = [[32, 16, 8, 4], [32, 16, 8, 4]]
        elif iteration == 'SyN':
            reg = ['Rigid', 'Affine', 'SyN']
            params = [(0.25,), (0.25,), (0.15, 3, 0.5)]
            sigma = [[4, 4, 2, 2], [4, 4, 2, 2], [4, 4, 4, 2]]
            shrink = [[32, 16, 8, 4], [32, 16, 8, 4], [64, 32, 16, 8]]
            if i == len(reg_types) - 1:
                sigma = [[4, 4, 2, 2], [4, 4, 2, 2], [4, 4, 4, 2, 1]]
                shrink = [[32, 16, 8, 4], [32, 16, 8, 4], [64, 32, 16, 8, 4]]
        else:
            sys.exit('The registration type "' + str(iteration) + '" does not exist. Reconstruction stopped.')

        if template_present:
            all_images.extend([template])

        for image in all_images:
            reg_output = cur_it_dir + image.split('/')[-1]
            registration(average, image, reg_output, reg, sigma, shrink, params, dim=3)

        registered_images = files_ind_dir_wo_string(cur_it_dir, '-Transformation')
        registered_images_wo_allen = []

        for reg_image in registered_images:
            if 'allen' in reg_image:
                registered_allen = reg_image
            else:
                registered_images_wo_allen.append(reg_image)

        robust_avg_volumes(registered_images_wo_allen, out_dir + str(i+1).zfill(2) + '-average.nii')
        if template_present:
            avg_template(out_dir + str(i+1).zfill(2) + '-average.nii',
                         registered_allen,
                         template_weight,
                         out_dir + str(i + 1).zfill(2) + '-average-template.nii')


if __name__ == '__main__':
    arguments = sys.argv

    input_symmetric = False
    input_robust = True
    input_template = ''
    input_weight = 1
    show_log = False

    if len(arguments) == 1:
        print(usage_message)
    else:
        try:
            opts, args = getopt.getopt(arguments[1:], "hi:o:s:t:w:l:", ["images=", "output=", "symmetric=",
                                                                        "template=", "weight=", "log="])
        except getopt.GetoptError:
            print('\nSomething is not right in your syntax. Tak a look at your possibilities:\n', usage_message)
            sys.exit(2)
        for opt, arg in opts:
            if opt == "-h":
                print(usage_message)
                sys.exit()
            elif opt in ("-i", "--images"):
                input_images = arg
            elif opt in ("-o", "--output"):
                output_directory = arg
            elif opt in ("-s", "--symmetric"):
                if arg == 'True' or arg == 'true':
                    input_symmetric = True
                elif arg == 'False' or arg == 'false':
                    input_symmetric = False
                else:
                    print('\nSomething is not right in your syntax. Tak a look at your possibilities:\n', usage_message)
            elif opt in ("-t", "--template"):
                input_template = arg
            elif opt in ("-w", "--weight"):
                input_weight = arg
            elif opt in ("-l", "--log"):
                if arg == 'True' or arg == 'true':
                    show_log = True
                elif arg == 'False' or arg == 'false':
                    show_log = False
                else:
                    print('\nSomething is not right in your syntax. Tak a look at your possibilities:\n', usage_message)
                    sys.exit(2)
        try:
            create_template_from_images(input_images,
                                        output_directory,
                                        symmetric=input_symmetric,
                                        template=input_template,
                                        template_weight=input_weight,
                                        print_log=show_log)
        except NameError:
            print('\nYou are not (properly) defining all the input variables. Tak a look at what is needed:\n',
                  usage_message)