-
Notifications
You must be signed in to change notification settings - Fork 33
/
test.py
108 lines (79 loc) · 3.37 KB
/
test.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
import os
import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
from pytorch_msssim import ssim
from torch.utils.data import DataLoader
from collections import OrderedDict
from utils import AverageMeter, write_img, chw_to_hwc
from datasets.loader import PairLoader
from models import *
parser = argparse.ArgumentParser()
parser.add_argument('--model', default='dehazeformer-s', type=str, help='model name')
parser.add_argument('--num_workers', default=16, type=int, help='number of workers')
parser.add_argument('--data_dir', default='./data/', type=str, help='path to dataset')
parser.add_argument('--save_dir', default='./saved_models/', type=str, help='path to models saving')
parser.add_argument('--result_dir', default='./results/', type=str, help='path to results saving')
parser.add_argument('--dataset', default='RESIDE-IN', type=str, help='dataset name')
parser.add_argument('--exp', default='indoor', type=str, help='experiment setting')
args = parser.parse_args()
def single(save_dir):
state_dict = torch.load(save_dir)['state_dict']
new_state_dict = OrderedDict()
for k, v in state_dict.items():
name = k[7:]
new_state_dict[name] = v
return new_state_dict
def test(test_loader, network, result_dir):
PSNR = AverageMeter()
SSIM = AverageMeter()
torch.cuda.empty_cache()
network.eval()
os.makedirs(os.path.join(result_dir, 'imgs'), exist_ok=True)
f_result = open(os.path.join(result_dir, 'results.csv'), 'w')
for idx, batch in enumerate(test_loader):
input = batch['source'].cuda()
target = batch['target'].cuda()
filename = batch['filename'][0]
with torch.no_grad():
output = network(input).clamp_(-1, 1)
# [-1, 1] to [0, 1]
output = output * 0.5 + 0.5
target = target * 0.5 + 0.5
psnr_val = 10 * torch.log10(1 / F.mse_loss(output, target)).item()
_, _, H, W = output.size()
down_ratio = max(1, round(min(H, W) / 256)) # Zhou Wang
ssim_val = ssim(F.adaptive_avg_pool2d(output, (int(H / down_ratio), int(W / down_ratio))),
F.adaptive_avg_pool2d(target, (int(H / down_ratio), int(W / down_ratio))),
data_range=1, size_average=False).item()
PSNR.update(psnr_val)
SSIM.update(ssim_val)
print('Test: [{0}]\t'
'PSNR: {psnr.val:.02f} ({psnr.avg:.02f})\t'
'SSIM: {ssim.val:.03f} ({ssim.avg:.03f})'
.format(idx, psnr=PSNR, ssim=SSIM))
f_result.write('%s,%.02f,%.03f\n'%(filename, psnr_val, ssim_val))
out_img = chw_to_hwc(output.detach().cpu().squeeze(0).numpy())
write_img(os.path.join(result_dir, 'imgs', filename), out_img)
f_result.close()
os.rename(os.path.join(result_dir, 'results.csv'),
os.path.join(result_dir, '%.02f | %.04f.csv'%(PSNR.avg, SSIM.avg)))
if __name__ == '__main__':
network = eval(args.model.replace('-', '_'))()
network.cuda()
saved_model_dir = os.path.join(args.save_dir, args.exp, args.model+'.pth')
if os.path.exists(saved_model_dir):
print('==> Start testing, current model name: ' + args.model)
network.load_state_dict(single(saved_model_dir))
else:
print('==> No existing trained model!')
exit(0)
dataset_dir = os.path.join(args.data_dir, args.dataset)
test_dataset = PairLoader(dataset_dir, 'test', 'test')
test_loader = DataLoader(test_dataset,
batch_size=1,
num_workers=args.num_workers,
pin_memory=True)
result_dir = os.path.join(args.result_dir, args.dataset, args.model)
test(test_loader, network, result_dir)