def generate_atlas_set(original_txt_path,atlas_path,l_atlas_path, output_path,phase='train',test_phase_path_list=None, test_phase_l_path_list=None):
if phase!="test":
source_path_list,target_path_list,l_source_path_list, l_target_path_list=loading_img_list_from_files(original_txt_path)
else:
source_path_list =test_phase_path_list
l_source_path_list = test_phase_l_path_list
target_path_list = []
l_target_path_list = []
source_path_list =source_path_list+target_path_list
file_num = len(source_path_list)
l_source_path_list = l_source_path_list+l_target_path_list
target_path_list = [atlas_path for _ in range(file_num)]
l_target_path_list = [l_atlas_path for _ in range(file_num)]
if l_source_path_list is not None and l_target_path_list is not None:
assert len(source_path_list) == len(l_source_path_list)
file_list = [[source_path_list[i], target_path_list[i],l_source_path_list[i],l_target_path_list[i]] for i in range(file_num)]
else:
file_list = [[source_path_list[i], target_path_list[i]] for i in range(file_num)]
output_phase_path = os.path.join(output_path,phase)
os.makedirs(output_phase_path,exist_ok=True)
pair_txt_path = os.path.join(output_phase_path,'pair_path_list.txt')
fn_txt_path = os.path.join(output_phase_path,'pair_name_list.txt')
fname_list = [generate_pair_name([file_list[i][0],file_list[i][1]]) for i in range(file_num)]
write_list_into_txt(pair_txt_path,file_list)
write_list_into_txt(fn_txt_path,fname_list)
def generate_moving_target_dict(pair_path_list_txt, pair_name_list_txt=None):
pair_path_list = read_img_label_into_list(pair_path_list_txt)
if pair_name_list_txt is None:
pair_name_detail_list = [
generate_pair_name([pair_path[0], pair_path[1]], detail=True)
for pair_path in pair_path_list
]
else:
pair_name_detail_list = read_fname_list_from_pair_fname_txt(
pair_name_list_txt, detail=True)
moving_name_list = [name_list[1] for name_list in pair_name_detail_list]
moving_name_set = set(moving_name_list)
moving_target_dict = {
moving_name: {
'm_pth': None,
"t_pth": [],
"name": [],
"l_pth": None
}
for moving_name in moving_name_set
}
for i in range(len(moving_name_list)):
has_label = len(pair_path_list[i]) == 4
moving_name = moving_name_list[i]
moving_target_dict[moving_name]['m_pth'] = pair_path_list[i][0]
moving_target_dict[moving_name]["t_pth"].append(pair_path_list[i][1])
moving_target_dict[moving_name]["name"].append(
pair_name_detail_list[i])
if has_label:
moving_target_dict[
moving_name_list[i]]["l_pth"] = pair_path_list[i][2]
return moving_target_dict
def gen_intra_pair_list(img_path_list,
fname_list,
label_path_list,
pair_num_limit=-1,
per_num_limit=-1):
"""
:param img_path_list: NxK list
:param label_path_list: NxK list
:param pair_num_limit: max number of the random pair
:param per_num_limit: max number of a image being the source image
:return:
"""
img_pair_list = []
pair_name_list = []
num_img = len(img_path_list)
for i in range(num_img):
img_pair_list_tmp = []
pair_name_list_tmp = []
for j in range(1, len(img_path_list[i])):
if label_path_list[i][j] is not None:
img_pair_list_tmp.append([
img_path_list[i][0], img_path_list[i][j],
label_path_list[i][0], label_path_list[i][j]
])
else:
img_pair_list_tmp.append(
[img_path_list[i][0], img_path_list[i][j]])
if fname_list is not None:
pair_name_list_tmp.append([
fname_list[i][0] + "_" + fname_list[i][j],
fname_list[i][0], fname_list[i][j]
])
else:
pair_name_list_tmp.append(
generate_pair_name(
[img_path_list[i][0], img_path_list[i][j]],
detail=True))
if len(img_pair_list_tmp) > per_num_limit and per_num_limit > -1:
ind = list(range(len(img_pair_list_tmp)))
random.shuffle(ind)
img_pair_list_tmp = [
img_pair_list_tmp[ind[i]] for i in range(per_num_limit)
]
pair_name_list_tmp = [
pair_name_list_tmp[ind[i]] for i in range(per_num_limit)
]
img_pair_list += img_pair_list_tmp
pair_name_list += pair_name_list_tmp
if len(img_pair_list) > pair_num_limit and pair_num_limit >= 0:
ind = list(range(len(img_pair_list)))
random.shuffle(ind)
img_pair_list = [img_pair_list[ind[i]] for i in range(pair_num_limit)]
pair_name_list = [
pair_name_list[ind[i]] for i in range(pair_num_limit)
]
return img_pair_list, pair_name_list
else:
return img_pair_list, pair_name_list
def split_input(original_txt_path):
source_path_list, target_path_list, l_source_path_list, l_target_path_list = loading_img_list_from_files(
original_txt_path)
file_num = len(source_path_list)
if l_source_path_list is not None and l_target_path_list is not None:
assert len(source_path_list) == len(l_source_path_list)
file_list = [[
source_path_list[i], target_path_list[i], l_source_path_list[i],
l_target_path_list[i]
] for i in range(file_num)]
else:
file_list = [[source_path_list[i], target_path_list[i]]
for i in range(file_num)]
fname_list = [
generate_pair_name([file_list[i][0], file_list[i][1]])
for i in range(file_num)
]
return file_list, fname_list
def init_env(output_path,
source_path_list,
target_path_list,
l_source_path_list=None,
l_target_path_list=None):
"""
:param task_full_path: the path of a completed task
:param source_path: path of the source image
:param target_path: path of the target image
:param l_source: path of the label of the source image
:param l_target: path of the label of the target image
:return: None
"""
file_num = len(source_path_list)
assert len(source_path_list) == len(target_path_list)
if l_source_path_list is not None and l_target_path_list is not None:
assert len(source_path_list) == len(l_source_path_list)
file_list = [[
source_path_list[i], target_path_list[i], l_source_path_list[i],
l_target_path_list[i]
] for i in range(file_num)]
else:
file_list = [[source_path_list[i], target_path_list[i]]
for i in range(file_num)]
os.makedirs(os.path.join(output_path, 'reg/test'), exist_ok=True)
os.makedirs(os.path.join(output_path, 'reg/res'), exist_ok=True)
pair_txt_path = os.path.join(output_path, 'reg/test/pair_path_list.txt')
fn_txt_path = os.path.join(output_path, 'reg/test/pair_name_list.txt')
fname_list = [
generate_pair_name([file_list[i][0], file_list[i][1]])
for i in range(file_num)
]
write_list_into_txt(pair_txt_path, file_list)
write_list_into_txt(fn_txt_path, fname_list)
root_path = output_path
data_task_name = 'reg'
cur_task_name = 'res'
return root_path
def init_test_env(setting_path,
output_path,
registration_pair_list,
pair_name_list=None):
"""
create test environment, the pair list would be saved into output_path/reg/test/pair_path_list.txt,
a corresponding auto-parsed filename list would also be saved in output/path/reg/test/pair_name_list.txt
:param setting_path: the path to load 'cur_task_setting.json' and 'cur_data_setting.json' (optional if the related settings are in cur_task_setting)
:param output_path: the output path of the task
:param registration_pair_list: including source_path_list, target_path_list, l_source_path_list, l_target_path_list
:return: tuple of ParameterDict, datapro (optional) and tsk_set
"""
source_path_list, target_path_list, l_source_path_list, l_target_path_list = registration_pair_list
dm_json_path = os.path.join(setting_path, 'cur_data_setting.json')
tsm_json_path = os.path.join(setting_path, 'cur_task_setting.json')
assert os.path.isfile(tsm_json_path), "task setting {} not exists".format(
tsm_json_path)
dm = DataTask('task_reg',
dm_json_path) if os.path.isfile(dm_json_path) else None
tsm = ModelTask('task_reg', tsm_json_path)
file_num = len(source_path_list)
if l_source_path_list is not None and l_target_path_list is not None:
file_list = [[
source_path_list[i], target_path_list[i], l_source_path_list[i],
l_target_path_list[i]
] for i in range(file_num)]
else:
file_list = [[source_path_list[i], target_path_list[i]]
for i in range(file_num)]
os.makedirs(os.path.join(output_path, 'reg/test'), exist_ok=True)
os.makedirs(os.path.join(output_path, 'reg/res'), exist_ok=True)
pair_txt_path = os.path.join(output_path, 'reg/test/pair_path_list.txt')
fn_txt_path = os.path.join(output_path, 'reg/test/pair_name_list.txt')
if pair_name_list is None:
pair_name_list = [
generate_pair_name([file_list[i][0], file_list[i][1]], detail=True)
for i in range(file_num)
]
write_list_into_txt(pair_txt_path, file_list)
write_list_into_txt(fn_txt_path, pair_name_list)
data_task_name = 'reg'
cur_task_name = 'res'
if dm is not None:
dm.data_par['datapro']['dataset']['output_path'] = output_path
dm.data_par['datapro']['dataset']['task_name'] = data_task_name
tsm.task_par['tsk_set']['task_name'] = cur_task_name
tsm.task_par['tsk_set']['output_root_path'] = os.path.join(
output_path, data_task_name)
if tsm.task_par['tsk_set']['model'] == 'reg_net':
tsm.task_par['tsk_set']['reg']['mermaid_net'][
'mermaid_net_json_pth'] = os.path.join(
setting_path, 'mermaid_nonp_settings.json')
if tsm.task_par['tsk_set']['model'] == 'mermaid_iter':
tsm.task_par['tsk_set']['reg']['mermaid_iter'][
'mermaid_affine_json'] = os.path.join(
setting_path, 'mermaid_affine_settings.json')
tsm.task_par['tsk_set']['reg']['mermaid_iter'][
'mermaid_nonp_json'] = os.path.join(setting_path,
'mermaid_nonp_settings.json')
return dm, tsm
def gen_post_aug_pair_list(test_img_path_list,
train_img_path_list,
test_fname_list=None,
train_fname_list=None,
test_label_path_list=None,
train_label_path_list=None,
pair_num_limit=-1,
per_num_limit=-1):
"""
:param test_img_path_list:
:param train_img_path_list:
:param test_fname_list:
:param train_fname_list:
:param test_label_path_list:
:param train_label_path_list:
:param pair_num_limit:
:param per_num_limit:
:return:
"""
img_pair_list = []
pair_name_list = []
num_test_img = len(test_img_path_list)
num_train_img = len(train_img_path_list)
if test_label_path_list is None:
test_label_path_list = ["None"] * num_test_img
if train_label_path_list is None:
train_label_path_list = ["None"] * num_train_img
for i in range(num_test_img):
img_pair_list_tmp = []
pair_name_list_tmp = []
for j in range(num_train_img):
img_pair_list_tmp.append([
test_img_path_list[i], train_img_path_list[j],
test_label_path_list[i], train_label_path_list[j]
])
if train_fname_list is not None and test_fname_list is not None:
pair_name_list_tmp.append([
test_fname_list[i] + "_" + train_fname_list[j],
test_fname_list[i], train_fname_list[j]
])
else:
pair_name_list_tmp.append(
generate_pair_name(
[test_img_path_list[i], train_img_path_list[j]],
detail=True))
if len(img_pair_list_tmp) > per_num_limit and per_num_limit > -1:
ind = list(range(len(img_pair_list_tmp)))
random.shuffle(ind)
img_pair_list_tmp = [
img_pair_list_tmp[ind[i]] for i in range(per_num_limit)
]
pair_name_list_tmp = [
pair_name_list_tmp[ind[i]] for i in range(per_num_limit)
]
img_pair_list += img_pair_list_tmp
pair_name_list += pair_name_list_tmp
if len(img_pair_list) > pair_num_limit and pair_num_limit >= 0:
ind = list(range(len(img_pair_list)))
random.shuffle(ind)
img_pair_list = [img_pair_list[ind[i]] for i in range(pair_num_limit)]
pair_name_list = [
pair_name_list[ind[i]] for i in range(pair_num_limit)
]
return img_pair_list, pair_name_list
else:
return img_pair_list, pair_name_list
output_path=out_path_list,
fixed_sz=desired_sz)
indexes = list(range(len(raw_path_list)))
number_of_workers = 8
boundary_list = []
config = dict()
index_partitions = np.array_split(indexes, number_of_workers)
if process_img:
with Pool(processes=number_of_workers) as pool:
pool.map(f, index_partitions)
if generate_txt:
source_list = [f['source'] for f in out_path_list]
target_list = [f['target'] for f in out_path_list]
lsource_list = [f['lsource'] for f in out_path_list]
ltarget_list = [f['ltarget'] for f in out_path_list]
file_num = len(source_list)
file_list = [[
source_list[i], target_list[i], lsource_list[i], ltarget_list[i]
] for i in range(file_num)]
os.makedirs(os.path.join(output_txt_path, 'reg/test'), exist_ok=True)
os.makedirs(os.path.join(output_txt_path, 'reg/res'), exist_ok=True)
pair_txt_path = os.path.join(output_txt_path,
'reg/test/pair_path_list.txt')
fn_txt_path = os.path.join(output_txt_path, 'reg/test/pair_name_list.txt')
fname_list = [
generate_pair_name([file_list[i][0], file_list[i][1]])
for i in range(file_num)
]
write_list_into_txt(pair_txt_path, file_list)
write_list_into_txt(fn_txt_path, fname_list)