-
Notifications
You must be signed in to change notification settings - Fork 2
/
Copy pathevaluate_model_group.py
186 lines (142 loc) · 9 KB
/
evaluate_model_group.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
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
# Script to evaluate a certain pix2pix hyperparameter
from general_imports import *
import argparse
import random
from utils_stereofog import calculate_model_results
parser = argparse.ArgumentParser()
parser.add_argument('--results_path', required=True, help='path to the results folder (e.g., results/hyperparameters/hyperparameter_netG)')
parser.add_argument('--ratio', type=float, default=4/3, help='aspect ratio of the images (to undo the transformation from pix2pix model)')
parser.add_argument('--num_images', type=int, default=5, help='number of images to plot for evaluation')
parser.add_argument('--shuffle', action='store_true', help='if specified, shuffle the images before plotting')
parser.add_argument('--seed', type=int, default=16, help='seed for the random shuffling of the images')
parser.add_argument('--flip', action='store_true', help='if specified, make the columns correspond to different test images, not different models (i.e., rows and columns are flipped)')
parser.add_argument('--epoch', type=int, default=-1, help='epoch to plot the results of')
parser.add_argument('--dont_save_figure', action='store_true', help='do not save the final evaluation figure')
results_path = parser.parse_args().results_path
ratio = parser.parse_args().ratio
num_images = parser.parse_args().num_images
shuffle = parser.parse_args().shuffle
seed = parser.parse_args().seed
flip = parser.parse_args().flip
epoch = parser.parse_args().epoch
dont_save_figure = parser.parse_args().dont_save_figure
fontsize = 20
if epoch == -1:
epoch = 'latest'
if 'final_models' in results_path:
hyperparameter = 'final_models'
else:
hyperparameter = results_path.split('/')[-1].replace('hyperparameter_', '')
subfolders = [item for item in os.listdir(results_path) if os.path.isdir(os.path.join(results_path, item))]
num_models = len(subfolders)
if results_path.split('/')[-1].count('_') == 1:
models = [item.replace(results_path.split('/')[-1] + '_', '') for item in subfolders]
else:
models = [item.replace('hyperparameter_', '') for item in subfolders]
both = sorted(zip(models, subfolders))
models = [model for model, subfolder in both]
subfolders = [subfolder for model, subfolder in both]
# Part that needs to be inserted due to the way pix2pix saves the images
subpath_addition = f'test_{epoch}/images'
images = [item for item in os.listdir(os.path.join(results_path, subfolders[0], subpath_addition)) if 'fake_B' in item]
if shuffle:
random.seed(seed)
random.shuffle(images)
try:
images_to_plot = images[:num_images]
except:
print(f'Number of images to plot ({num_images}) is greater than the number of images available ({len(images)}). Defaulting to {len(images)} images.')
images_to_plot = images
limit = len(images_to_plot)
# os.listdir('./datasets/stereofog_data/A/test')[0]
# orig_img = plt.imread('./datasets/stereofog_data/A/test/2023-08-01-04__20.bmp')
# ratio = orig_img.shape[1]/orig_img.shape[0]
width_per_image = 4
height_per_image = width_per_image / ratio
# ax = [fig.add_subplot(num_models+2,limit,i+1) for i in range(limit*(num_models+2))]
# ax[0].text(0.5, 1.05, 'fake', fontsize=15, color='k', fontweight='black', ha='center', transform=ax[0].transAxes)
# ax[1].text(0.5, 1.05, 'foggy real', fontsize=15, color='k', fontweight='black', ha='center', transform=ax[1].transAxes)
# ax[2].text(0.5, 1.05, 'clear real', fontsize=15, color='k', fontweight='black', ha='center', transform=ax[2].transAxes)
# Calculating evaluation scores for each model
Pearson_correlation_scores = []
MSE_scores = []
PSNR_scores = []
NCC_scores = []
SSIM_scores = []
CW_SSIM_scores = []
MS_SSIM_scores = []
for model_index, model in enumerate(models):
mean_Pearson, mean_MSE, mean_PSNR, mean_NCC, mean_SSIM, mean_CW_SSIM, mean_MS_SSIM = calculate_model_results(os.path.join(results_path, subfolders[model_index]), epoch=epoch)
Pearson_correlation_scores.append(mean_Pearson)
MSE_scores.append(mean_MSE)
PSNR_scores.append(mean_PSNR)
NCC_scores.append(mean_NCC)
SSIM_scores.append(mean_SSIM)
CW_SSIM_scores.append(mean_CW_SSIM)
MS_SSIM_scores.append(mean_MS_SSIM)
scores = { 'Pearson': Pearson_correlation_scores,
'MSE': MSE_scores,
'PSNR': PSNR_scores,
'NCC': NCC_scores,
'SSIM': SSIM_scores,
'CW-SSIM': CW_SSIM_scores,
'MS-SSIM': MS_SSIM_scores
}
if not flip:
fig, ax = plt.subplots(limit, num_models+2, figsize=(width_per_image*(num_models+2), height_per_image*limit))
ax[0, 0].text(0.5,1.105, 'foggy real', transform=ax[0, 0].transAxes, backgroundcolor='w', horizontalalignment='center', verticalalignment='center', fontsize=fontsize, fontweight='black', color='k')
ax[0, 1].text(0.5,1.105, 'ground truth', transform=ax[0, 1].transAxes, backgroundcolor='w', horizontalalignment='center', verticalalignment='center', fontsize=fontsize, fontweight='black', color='k')
for i in range(limit):
# Plotting the ground truth image
img_original = plt.imread(os.path.join(results_path, subfolders[0], subpath_addition, images_to_plot[i][:-10] + 'real_A' + '.png'))
ax[i, 0].imshow(img_original, aspect='auto')
ax[i, 0].axis('off')
# Plotting the fogged image
img_fogged = plt.imread(os.path.join(results_path, subfolders[0], subpath_addition, images_to_plot[i][:-10] + 'real_B' + '.png'))
ax[i, 1].imshow(img_fogged, aspect='auto')
ax[i, 1].axis('off')
# Plotting each of the model's results
for j in range(num_models):
img = plt.imread(os.path.join(results_path, subfolders[j], subpath_addition, images_to_plot[i]))
ax[i, j+2].imshow(img, aspect='auto')
if i == 0:
ax[i, j+2].text(0.5,1.105, models[j], transform=ax[i, j+2].transAxes, horizontalalignment='center', verticalalignment='center', fontsize=fontsize, fontweight='black', color='k', zorder=20)
ax[i, j+2].axis('off')
# Plotting each model's scores
score_plot_distance = 0.1
for j in range(num_models):
for score_index, score in enumerate(scores.keys()):
ax[0, j+2].text(0.5,1.2 + score_plot_distance*score_index, f'{score}: {scores[score][j]:.2f}', transform=ax[0, j+2].transAxes, horizontalalignment='center', verticalalignment='center', fontsize=14, fontweight='black', color='k')
else:
fig, ax = plt.subplots(num_models+2, limit, figsize=(limit*width_per_image,(num_models+2)*height_per_image)) # num_models+2 to acommodate ground truth and fogged image
ax[0, 0].text(-0.07,0.5, 'foggy real', transform=ax[0, 0].transAxes, backgroundcolor='w', horizontalalignment='center', verticalalignment='center', fontsize=fontsize, fontweight='black', color='k', rotation='vertical')
ax[1, 0].text(-0.07,0.5, 'ground truth', transform=ax[1, 0].transAxes, backgroundcolor='w', horizontalalignment='center', verticalalignment='center', fontsize=fontsize, fontweight='black', color='k', rotation='vertical')
for i in range(limit):
# Plotting the ground truth image
img_original = plt.imread(os.path.join(results_path, subfolders[0], subpath_addition, images_to_plot[i][:-10] + 'real_A' + '.png'))
ax[0, i].imshow(img_original, aspect='auto')
ax[0, i].axis('off')
# Plotting the fogged image
img_fogged = plt.imread(os.path.join(results_path, subfolders[0], subpath_addition, images_to_plot[i][:-10] + 'real_B' + '.png'))
ax[1, i].imshow(img_fogged, aspect='auto')
ax[1, i].axis('off')
# Plotting each of the model's results
for j in range(num_models):
img = plt.imread(os.path.join(results_path, subfolders[j], subpath_addition, images_to_plot[i]))
ax[j+2, i].imshow(img, aspect='auto')
if i == 0:
ax[j+2, i].text(-0.07,0.5, models[j], transform=ax[j+2, i].transAxes, backgroundcolor='w', horizontalalignment='center', verticalalignment='center', fontsize=fontsize, fontweight='black', color='k', rotation='vertical', zorder=20)
ax[j+2, i].axis('off')
# Plotting each model's scores
score_plot_distance = 0.1
for j in range(num_models):
for score_index, score in enumerate(scores.keys()):
ax[j+2, 0].text(-0.15 - score_plot_distance*score_index,0.5, f'{score}: {scores[score][j]:.2f}', transform=ax[j+2, 0].transAxes, backgroundcolor='w', horizontalalignment='center', verticalalignment='center', fontsize=12, fontweight='black', color='k', rotation='vertical')
# for j in range(num_models):
# ax[j+2, 0].text(-0.1,0.5, f'SSIM:{SSIM_scores[j]:.2f}\nCW-SSIM:{CW_SSIM_scores[j]:.2f}', transform=ax[j+2, 0].transAxes, backgroundcolor='w', horizontalalignment='center', verticalalignment='center', fontsize=12, fontweight='black', color='k', rotation='vertical')
plt.subplots_adjust(hspace=0, wspace=0)
if dont_save_figure:
plt.show()
else:
plt.savefig(os.path.join(f"{results_path}", f"{hyperparameter}_evaluation.pdf"), bbox_inches='tight', format='pdf', pad_inches=0)
print(f"Saved evaluation figure to {os.path.join(f'{results_path}', f'{hyperparameter}_evaluation.pdf')}.")