-
Notifications
You must be signed in to change notification settings - Fork 17
/
Copy pathAll_in_One_Multi_Process.py
135 lines (103 loc) · 5.41 KB
/
All_in_One_Multi_Process.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
import os
from tqdm import tqdm
import sys
import multiprocessing
import time
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
ir_dir = os.path.join(os.getcwd(), 'datasets.llvip/test_imgs/ir')
vi_dir = os.path.join(os.getcwd(), 'datasets.llvip/test_imgs/vi')
model_path_dict = dict()
model_path_dict_1 = dict()
model_path_dict_2 = dict()
model_path_dict_1['CSF'] = os.path.join(os.getcwd(), 'Checkpoint/CSF/EC.ckpt')
model_path_dict_2['CSF'] = os.path.join(os.getcwd(), 'Checkpoint/CSF/ED.ckpt')
model_path_dict_1['CUFD'] = os.path.join(os.getcwd(), 'Checkpoint/CUFD/1part1_model.ckpt')
model_path_dict_2['CUFD'] = os.path.join(os.getcwd(), 'Checkpoint/CUFD/part2_model.ckpt')
model_path_dict_1['DIDFuse'] = os.path.join(os.getcwd(), 'Checkpoint/DIDFuse/Encoder.pkl')
model_path_dict_2['DIDFuse'] = os.path.join(os.getcwd(), 'Checkpoint/DIDFuse/Decoder.pkl')
model_path_dict_1['DIVFusion'] = os.path.join(os.getcwd(), 'Checkpoint/DIVFusion/decom.ckpt')
model_path_dict_2['DIVFusion'] = os.path.join(os.getcwd(), 'Checkpoint/DIVFusion/enhance.ckpt')
model_path_dict_1['RFN-Nest'] = os.path.join(os.getcwd(), 'Checkpoint/RFN-Nest/RFN_Nest.model')
model_path_dict_2['RFN-Nest'] = os.path.join(os.getcwd(), 'Checkpoint/RFN-Nest/NestFuse.model')
model_path_dict['DenseFuse'] = os.path.join(os.getcwd(), 'Checkpoint/DenseFuse/DeseFuse.ckpt')
model_path_dict['FusionGAN'] = os.path.join(os.getcwd(), 'Checkpoint/FusionGAN/FusionGAN')
model_path_dict['GAN-FM'] = os.path.join(os.getcwd(), 'Checkpoint/GAN-FM/model.ckpt')
model_path_dict['GANMcC'] = os.path.join(os.getcwd(), 'Checkpoint/GANMcC/GANMcC')
model_path_dict['NestFuse'] = os.path.join(os.getcwd(), 'Checkpoint/NestFuse/nestfuse.model')
model_path_dict['PIAFusion'] = os.path.join(os.getcwd(), 'Checkpoint/PIAFusion')
model_path_dict['PMGI'] = os.path.join(os.getcwd(), 'Checkpoint/PMGI/PMGI')
model_path_dict['SDNet'] = os.path.join(os.getcwd(), 'Checkpoint/SDNet/SDNet.model')
model_path_dict['STDFusionNet'] = os.path.join(os.getcwd(), 'Checkpoint/STDFusionNet/Fusion.model-29')
model_path_dict['SeAFusion'] = os.path.join(os.getcwd(), 'Checkpoint/SeAFusion/SeAFusion.pth')
model_path_dict['SuperFusion'] = os.path.join(os.getcwd(), 'Checkpoint/SuperFusion/MSRS.pth')
model_path_dict['SwinFusion'] = os.path.join(os.getcwd(), 'Checkpoint/SwinFusion/SwinFusion.pth')
model_path_dict['TarDAL'] = os.path.join(os.getcwd(), 'Checkpoint/TarDAL/tardal++.pt')
model_path_dict['U2Fusion'] = os.path.join(os.getcwd(), 'Checkpoint/U2Fusion/model.ckpt')
model_path_dict['IFCNN'] = os.path.join(os.getcwd(), 'Checkpoint/IFCNN/IFCNN-MAX.pth')
model_path_dict['UMF-CMGR'] = os.path.join(os.getcwd(), 'Checkpoint/UMF-CMGR/UMF_CMGR.pth')
Method_list = [
'CSF', 'CUFD', 'DIDFuse', 'DIVFusion', 'DenseFuse',
'FusionGAN', 'GAN-FM', 'GANMcC', 'IFCNN', 'NestFuse',
'PIAFusion', 'PMGI', 'RFN-Nest', 'SDNet', 'STDFusionNet',
'SeAFusion', 'SuperFusion', 'SwinFusion', 'TarDAL', 'U2Fusion',
'UMF-CMGR'
]
print(len(Method_list))
two_model_list = ['CSF', 'CUFD', 'DIDFuse', 'DIVFusion', 'RFN-Nest']
for Method in Method_list:
save_dir = os.path.join(os.getcwd(), 'Results.llvip/', Method)
if Method not in two_model_list:
with open('script_' + Method + '.sh', 'w') as f:
f.write('#!/bin/bash\n')
f.write("cd {}\n".format(Method))
print(Method.replace('-', ''))
f.write("CUDA_VISIBLE_DEVICES=0 \
python {}.py \
--Method {} \
--model_path {} \
--ir_dir {}\
--vi_dir {} \
--save_dir {} \
--is_RGB {}\n".format(Method.replace('-', ''), Method, model_path_dict[Method], ir_dir, vi_dir, save_dir, True))
f.write("cd ..\n".format(Method))
# os.system('bash script.sh')
else:
with open('script_' + Method + '.sh', 'w') as f:
f.write('#!/bin/bash\n')
f.write("cd {}\n".format(Method))
print(Method.replace('-', ''))
f.write("CUDA_VISIBLE_DEVICES=0 \
python {}.py \
--Method {} \
--model_path_1 {} \
--model_path_2 {} \
--ir_dir {}\
--vi_dir {} \
--save_dir {} \
--is_RGB {}\n".format(Method.replace('-', ''), Method, model_path_dict_1[Method], model_path_dict_2[Method], ir_dir, vi_dir, save_dir, True))
f.write("cd ..\n".format(Method))
# os.system('bash script.sh')
# multi process func
def multi_process(m_list):
print('---', m_list)
for m in m_list:
print('bash' + ' ' + 'script_' + m + '.sh')
os.system('bash' + ' ' + 'script_' + m + '.sh')
# set groups
num_per_group = 4
Method_len = len(Method_list)
if Method_len % num_per_group == 0:
groups = Method_len // num_per_group
else:
groups = Method_len // num_per_group + 1
print("set groups, len: {}, groups: {}, num_per_group: {}". format(Method_len, groups, num_per_group))
# for i in range(0, Method_len, num_per_group):
# print(Method_list[i:i+num_per_group])
print('multi process...')
pool = multiprocessing.Pool(processes=groups)
for i in range(0, Method_len, num_per_group):
# print("---> ", Method_list[i:i+num_per_group])
pool.apply_async(multi_process, args=(Method_list[i:i+num_per_group], ))
pool.close()
pool.join()
print("multi core end.")