--- a +++ b/create_script.py @@ -0,0 +1,145 @@ +import sys +import argparse +import os + +from utils.logger import logger +from utils.filemanager import create_directory_if_not_exists, get_paths, check_paths, extract_parameter + +if __name__ == "__main__": + # optional arguments from the command line + parser = argparse.ArgumentParser() + + parser.add_argument('--dataset_path', type=str, default='dataset/train', help='root dir for nifti data') + parser.add_argument('--experiment_name', type=str, default='elastix_01', help='experiment name') + parser.add_argument('--parameters_path', type=str, default='elastix-parameters/Par0003', help='root dir for elastix parameters. The script will use all the parameters in this directory. A single .txt file can also be used.') + parser.add_argument('--output_path', type=str, default='output', help='root dir for output scripts') + parser.add_argument("--use_masks", action='store_true', help='if True, segmentation masks will be used during the registration.') + + # parse the arguments + args = parser.parse_args() + + # check if parameters_path is .txt file + if os.path.isfile(args.parameters_path): + parameters = args.parameters_path + + # cteate params key folder name + reg_params = '-p "{}"'.format(parameters).replace('\\', '/') + reg_params_key = parameters.split('/')[-1].replace('.txt', '') + transform_idx = 0 + + elif os.path.isdir(args.parameters_path): + # get the parameters from the parameters_path + parameters = os.listdir(args.parameters_path) + + if len(parameters) == 0: + logger.error(f"No parameters found in {args.parameters_path} directory.") + sys.exit(1) + + # cteate params key folder name + reg_params = ' '.join(['-p "{}"'.format(os.path.join(args.parameters_path, param)) for param in parameters]).replace('\\', '/') + reg_params_key = '+'.join(['{}'.format(param.replace('.txt', '')) for param in parameters]) + transform_idx = len(parameters) - 1 + + # get the name of the parameters folder + params_folder_name = extract_parameter(args.parameters_path) # Par0003, Par0004, ... + + # create experiment output + # args.experiment_name is useful to distinguish between different experiments (e.g. with or without preprocessing, etc.) + # reg_params_key is useful to distinguish between different registration parameters (e.g. Par0003, Par0004, etc.) + args.exp_output = os.path.join(args.output_path, args.experiment_name, reg_params_key) + create_directory_if_not_exists(args.exp_output) + + # get the exhale and inhale volumes and segmentations + exhale_volumes = get_paths(args, "eBHCT") + inhale_volumes = get_paths(args, "iBHCT") + + check_paths(args, exhale_volumes, "exhale volumes") + check_paths(args, inhale_volumes, "inhale volumes") + + # get the exhale and inhale segmentations if args.use_masks is True + exhale_seg = get_paths(args, "eBHCT_lung") if args.use_masks else [0 for _ in range(len(exhale_volumes))] # the list has to have values for the zip(*) to return the values inside + inhale_seg = get_paths(args, "iBHCT_lung") if args.use_masks else [0 for _ in range(len(inhale_volumes))] + + if args.use_masks: + check_paths(args, exhale_seg, "exhale segmentations") + check_paths(args, inhale_seg, "inhale segmentations") + + logger.info(f"Creating .bat file in {args.exp_output} output directory for {args.experiment_name} experiment.") + logger.info(f"Experimenting using {reg_params} params command...") + logger.info(f"Key for the experiment: {reg_params_key}...") + + # creating the .bat file + with open(os.path.join(args.exp_output, 'elastix_transformix.bat'), 'w') as file: + file.write("@echo on\n") + file.write(f"echo To execute this file, use: call {os.path.join(args.exp_output, 'elastix_transformix.bat')} \n") + + for e_path, i_path, e_seg_path, i_seg_path in zip(exhale_volumes, inhale_volumes, exhale_seg, inhale_seg): + # Append commands to the .bat file + file.write(f"\nREM Processing {e_path} and {i_path}\n") + + # defining the sample name and the file names + sample_name = i_path.split('/')[-1].split('_')[0] #copd1, copd2, ... + e_filename_full = e_path.split('/')[-1].split('.')[0] #copd1_eBHCT, .. + i_filename_full = i_path.split('/')[-1].split('.')[0] #copd1_iBHCT, .. + + # defining fixed and moving paths + fixed_path = i_path + fMask = i_seg_path + + moving_path = e_path + mMask = e_seg_path + + # defining the control point to be transformed abd transform file paths + input_points = os.path.join(args.dataset_path ,f'{sample_name}/{sample_name}_300_iBH_xyz_r1.txt') # i_cntl_pt + transform_path = f'{args.exp_output}/images/output_{i_filename_full}/{e_filename_full}/TransformParameters.{transform_idx}.txt' + + # Get the names of the fixed and moving images for the output directory, names without the file extensions + reg_fixed_name = fixed_path.replace("\\", "/").split("/")[-1].split(".")[0] + reg_moving_name = moving_path.replace("\\", "/").split("/")[-1].split(".")[0] + + elastix_output_dir = f'{args.exp_output}/images/output_{reg_fixed_name}/{reg_moving_name}' + transformix_output_dir = f'{args.exp_output}/points/output_{reg_fixed_name}/{reg_moving_name}' + + create_directory_if_not_exists(elastix_output_dir) + create_directory_if_not_exists(transformix_output_dir) + + if params_folder_name == 'Par0003': + # elastix version: 3.9 -- fails to allocate memory to register; also can't -def point transformation + # elastix_v_path = '.\\elastix-versions\\elastix_windows32_v3.9\\elastix' + # transformix_v_path = '.\\elastix-versions\\elastix_windows32_v3.9\\transformix' + + # using default is 4.7 + elastix_v_path = '.\\elastix-versions\\elastix_windows64_v4.7\\elastix' + transformix_v_path = '.\\elastix-versions\\elastix_windows64_v4.7\\transformix' + + elif params_folder_name == 'Par0007': + # elastix version: 4.0 -- fails to allocate memory to register; also can't -def point transformation + # elastix_v_path = '.\\elastix-versions\\elastix_windows32_v4.0\\elastix' + # transformix_v_path = '.\\elastix-versions\\elastix_windows32_v4.0\\transformix' + + # using default is 4.7 + elastix_v_path = '.\\elastix-versions\\elastix_windows64_v4.7\\elastix' + transformix_v_path = '.\\elastix-versions\\elastix_windows64_v4.7\\transformix' + + elif params_folder_name == 'Par0011': + # elastix version: 4.301 + elastix_v_path = '.\\elastix-versions\\elastix_windows64_v4.3\\elastix' + transformix_v_path = '.\\elastix-versions\\elastix_windows64_v4.3\\transformix' + else: + # default is 4.7 + elastix_v_path = '.\\elastix-versions\\elastix_windows64_v4.7\\elastix' + transformix_v_path = '.\\elastix-versions\\elastix_windows64_v4.7\\transformix' + + + # create elastix command line + if args.use_masks: + elastix_command_line = f'{elastix_v_path} -f "{fixed_path}" -m "{moving_path}" -fMask "{fMask}" -mMask "{mMask}" {reg_params} -out "{elastix_output_dir}"' + else: + elastix_command_line = f'{elastix_v_path} -f "{fixed_path}" -m "{moving_path}" {reg_params} -out "{elastix_output_dir}"' + + # create transformix command line + trasformix_command_line = f'{transformix_v_path} -def "{input_points}" -tp "{transform_path}" -out "{transformix_output_dir}"' + + file.write(f"{elastix_command_line}\n") + file.write(f"{trasformix_command_line}\n") + \ No newline at end of file