-
Notifications
You must be signed in to change notification settings - Fork 0
/
create_script.py
145 lines (110 loc) · 7.82 KB
/
create_script.py
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
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
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")