|
@@ -0,0 +1,288 @@
|
|
|
+from helper_functions import pretty, make_dirs, slice_volume, files_in_dir, files_ind_dir_wo_string, make_volume
|
|
|
+from nipype.interfaces import ants
|
|
|
+from os import path
|
|
|
+from shutil import rmtree
|
|
|
+import sys
|
|
|
+import getopt
|
|
|
+
|
|
|
+# Help message.
|
|
|
+tab = ' '
|
|
|
+usage_message = '\nCOMMAND:\n' + tab + sys.argv[0] + \
|
|
|
+ '\n\nOPTIONS:\n' + \
|
|
|
+ tab + '-h, help menu\n' + \
|
|
|
+ tab + '-e, --template, <input template>\n' + \
|
|
|
+ tab + '-s, --slices, <directory with input slices>\n' + \
|
|
|
+ tab + '-r, --rawvolume, <raw volume of slices>\n' + \
|
|
|
+ tab + '-t, --inittrans, <initial transform, roughly transforming raw slice volume to match input ' \
|
|
|
+ 'volume, manually created>\n' + \
|
|
|
+ tab + '-d, --id, <brain id>\n' + \
|
|
|
+ tab + '-o, --output, <output directory, will create new directory if not present>\n' + \
|
|
|
+ tab + '-l, --log, [bool]>\n'
|
|
|
+
|
|
|
+
|
|
|
+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=['MeanSquares']*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,
|
|
|
+ write_composite_transform=True,
|
|
|
+ output_warped_image=result,
|
|
|
+ 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 co_register(template, slices, raw_volume, initial_transform, out_dir, brain_id, r=3, a=2, nl=3, print_log=False):
|
|
|
+ """
|
|
|
+ Function for reconstructing brain volume by iteratively registering to given template. The slices are stacked into a
|
|
|
+ 3-dimensional volume that is registered to the template. The template is inversely transformed and sliced, resulting
|
|
|
+ in slice pairs with a template slice and an input slice. For each pair, the input slice is registered to the
|
|
|
+ template slice, and all the resulting registered slices are stacked to a volume. This process is repeated r times
|
|
|
+ rigid, a times rigid and affine, and nl times rigid, affine, and non-linear. Returns the path to the final
|
|
|
+ reconstructed brain volume.
|
|
|
+ :param template: Path to template used as a reference for brain reconstruction.
|
|
|
+ :param slices: List of paths to slices that needs to be reconstructed.
|
|
|
+ :param raw_volume: Path to the initial stacked volume, needed for first rigid registration step.
|
|
|
+ :param initial_transform:Initial 3-dimensional transform of the raw volume to the template. If empty string given,
|
|
|
+ they will be matched by the center of mass.
|
|
|
+ :param out_dir: Output directory.
|
|
|
+ :param brain_id: Mouse id.
|
|
|
+ :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: Return path to the final reconstructed brain volume.
|
|
|
+ """
|
|
|
+
|
|
|
+ # Determine whether the reference template is the Allen atlas (1. iteration) or the DAPI-template (2. iteration).
|
|
|
+ if 'allen' in template:
|
|
|
+ template_type = 'allen'
|
|
|
+ else:
|
|
|
+ template_type = 'dapi'
|
|
|
+ volume_postfix = '_hv_reg_' + template_type + '_' + brain_id
|
|
|
+
|
|
|
+ # List of slices.
|
|
|
+ raw_slices = files_in_dir(slices, '.nii')
|
|
|
+
|
|
|
+ if print_log:
|
|
|
+ print('Volume given interpreted as:', template_type)
|
|
|
+ print('Checking paths.')
|
|
|
+
|
|
|
+ if not path.exists(raw_volume):
|
|
|
+ print('The path:', raw_volume, 'does not exist.')
|
|
|
+ if not path.exists(template):
|
|
|
+ print('The path:', template, 'does not exist.')
|
|
|
+ if not path.exists(initial_transform):
|
|
|
+ print('The path:', initial_transform, 'does not exist.')
|
|
|
+
|
|
|
+ if print_log:
|
|
|
+ print('Defining iterations.')
|
|
|
+ reg_types = ['Rigid'] * r + ['Affine'] * a + ['SyN'] * nl
|
|
|
+
|
|
|
+ if print_log:
|
|
|
+ print(reg_types)
|
|
|
+
|
|
|
+ # Run through iterations.
|
|
|
+ for i, iteration in enumerate(reg_types):
|
|
|
+ cur_work_dir = out_dir + pretty(i) + '/'
|
|
|
+ make_dirs(cur_work_dir)
|
|
|
+
|
|
|
+ if print_log:
|
|
|
+ print('Running iteration', pretty(i))
|
|
|
+ print('Current work dir:', cur_work_dir)
|
|
|
+
|
|
|
+ if i == 0:
|
|
|
+ hv = raw_volume
|
|
|
+ else:
|
|
|
+ hv = out_dir + pretty(i) + volume_postfix + '.nii'
|
|
|
+ if print_log:
|
|
|
+ print('Deleting prev. iteration work dir:', out_dir + pretty(i - 1) + '/')
|
|
|
+ rmtree(out_dir + pretty(i - 1) + '/')
|
|
|
+
|
|
|
+ # Iteration specific registration parameters.
|
|
|
+ if iteration == 'Rigid':
|
|
|
+ reg = ['Rigid']
|
|
|
+ params = [(0.25,)]
|
|
|
+ vol_sigma = [[4, 4, 2, 2]]
|
|
|
+ vol_shrink = [[32, 16, 8, 4]]
|
|
|
+ slice_sigma = [[4, 4, 2, 2]]
|
|
|
+ slice_shrink = [[32, 16, 8, 4]]
|
|
|
+ if i == 0:
|
|
|
+ vol_sigma = [[4, 4, 4]]
|
|
|
+ vol_shrink = [[32, 16, 8]]
|
|
|
+ slice_sigma = [[4, 4, 2]]
|
|
|
+ slice_shrink = [[32, 16, 8]]
|
|
|
+ elif iteration == 'Affine':
|
|
|
+ reg = ['Rigid', 'Affine']
|
|
|
+ params = [(0.25,), (0.25,), ]
|
|
|
+ vol_sigma = [[4, 4, 2, 2], [4, 4, 2, 2]]
|
|
|
+ vol_shrink = [[32, 16, 8, 4], [32, 16, 8, 4]]
|
|
|
+ slice_sigma = [[4, 4, 2, 2], [4, 4, 2, 2]]
|
|
|
+ slice_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)]
|
|
|
+ vol_sigma = [[4, 4, 2, 2], [4, 4, 2, 2], [4, 4, 4, 2]]
|
|
|
+ vol_shrink = [[32, 16, 8, 4], [32, 16, 8, 4], [64, 32, 16, 8]]
|
|
|
+ slice_sigma = [[4, 4, 2, 2], [4, 4, 2, 2], [4, 4, 4, 2, 2]]
|
|
|
+ slice_shrink = [[32, 16, 8, 4], [32, 16, 8, 4], [64, 32, 16, 8, 4]]
|
|
|
+ if i == len(reg_types)-1:
|
|
|
+ slice_sigma = [[4, 4, 2, 2], [4, 4, 2, 2], [4, 4, 4, 2, 2, 1]]
|
|
|
+ slice_shrink = [[32, 16, 8, 4], [32, 16, 8, 4], [64, 32, 16, 8, 4, 2]]
|
|
|
+ else:
|
|
|
+ sys.exit('The registration type "' + str(iteration) + '" does not exist. Reconstruction stopped.')
|
|
|
+
|
|
|
+ # Define name and path for registered template to hv.
|
|
|
+ cur_reg_name = pretty(i) + '-' + template_type + '_reg_to_hv.nii'
|
|
|
+ cur_reg_path = cur_work_dir + cur_reg_name
|
|
|
+
|
|
|
+ if print_log:
|
|
|
+ print('Registering hv to given volume - either MR or atlas')
|
|
|
+
|
|
|
+ if not path.exists(cur_reg_path):
|
|
|
+ registration(template, hv, cur_work_dir + 'hv_to_template.nii', reg, vol_sigma,
|
|
|
+ vol_shrink, params, dim=3, initial_trans=initial_transform,
|
|
|
+ inverse_result=cur_reg_path)
|
|
|
+
|
|
|
+ # Define and create directory for sliced template.
|
|
|
+ template_slice_dir = cur_work_dir + cur_reg_name.split('.')[0] + '-slices/'
|
|
|
+ make_dirs(template_slice_dir)
|
|
|
+
|
|
|
+ if print_log:
|
|
|
+ print('Slicing', cur_reg_path, 'into', template_slice_dir)
|
|
|
+
|
|
|
+ # Slice template.
|
|
|
+ slice_volume(cur_reg_path, template_slice_dir)
|
|
|
+ template_slices = files_in_dir(template_slice_dir, '.nii')
|
|
|
+
|
|
|
+ if len(raw_slices) == len(template_slices):
|
|
|
+ print('Number of images matches.')
|
|
|
+ else:
|
|
|
+ sys.exit('Mismatch between number of images.')
|
|
|
+
|
|
|
+ # Make directory for slice to slice registrations.
|
|
|
+ reg_slices_dir = cur_work_dir + pretty(i) + '-raw_reg_template/'
|
|
|
+ make_dirs(reg_slices_dir)
|
|
|
+
|
|
|
+ # Registered slice postfix.
|
|
|
+ reg_slice_postfix = '_reg'
|
|
|
+
|
|
|
+ # Perform slice to slice registrations.
|
|
|
+ for n, (raw_slice, template_slice) in enumerate(zip(raw_slices, template_slices)):
|
|
|
+ reg_slice = reg_slices_dir + pretty(i) + '_S' + pretty(n) + reg_slice_postfix + '.nii'
|
|
|
+ if not path.exists(reg_slice):
|
|
|
+ registration(template_slice, raw_slice, reg_slice, reg,
|
|
|
+ slice_sigma, slice_shrink, params, 2)
|
|
|
+
|
|
|
+ # Stack registered slices into new volume.
|
|
|
+ reg_slices = files_ind_dir_wo_string(reg_slices_dir, exclusion_string='-Transformation')
|
|
|
+ final_volume = out_dir + pretty(i + 1) + volume_postfix + '.nii'
|
|
|
+ make_volume(reg_slices, final_volume)
|
|
|
+
|
|
|
+ # Return final reconstructed brain volume.
|
|
|
+ return final_volume
|
|
|
+
|
|
|
+
|
|
|
+if __name__ == '__main__':
|
|
|
+ arguments = sys.argv
|
|
|
+ log = False
|
|
|
+
|
|
|
+ if len(arguments) == 1:
|
|
|
+ print(usage_message)
|
|
|
+ else:
|
|
|
+ try:
|
|
|
+ opts, args = getopt.getopt(arguments[1:],
|
|
|
+ "he:s:r:t:d:o:l:",
|
|
|
+ ["template=", "slices=", "rawvolume", "inittrans", "id", "output", "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 ("-e", "--template"):
|
|
|
+ input_template = arg
|
|
|
+ elif opt in ("-s", "--slices"):
|
|
|
+ input_slices = arg
|
|
|
+ elif opt in ("-r", "--rawvolume"):
|
|
|
+ input_raw_volume = arg
|
|
|
+ elif opt in ("-t", "--inittrans"):
|
|
|
+ input_init_trans = arg
|
|
|
+ elif opt in ("-d", "--id"):
|
|
|
+ input_brain_id = arg
|
|
|
+ elif opt in ("-o", "--output"):
|
|
|
+ output_dir = arg
|
|
|
+ elif opt in ("-l", "--log"):
|
|
|
+ if arg == 'True' or arg == 'true':
|
|
|
+ input_print_log = True
|
|
|
+ elif arg == 'False' or arg == 'false':
|
|
|
+ input_print_log = False
|
|
|
+ else:
|
|
|
+ print('\nSomething is not right in your syntax. Tak a look at your possibilities:\n', usage_message)
|
|
|
+ sys.exit(2)
|
|
|
+ try:
|
|
|
+ co_register(input_template,
|
|
|
+ input_slices,
|
|
|
+ input_raw_volume,
|
|
|
+ input_init_trans,
|
|
|
+ output_dir,
|
|
|
+ input_brain_id,
|
|
|
+ print_log=input_print_log)
|
|
|
+ except NameError:
|
|
|
+ print('\nYou are not (properly) defining all the input variables. Tak a look at what is needed:\n',
|
|
|
+ usage_message)
|