forked from DengPingFan/SINet
-
Notifications
You must be signed in to change notification settings - Fork 0
/
MyTest.py
53 lines (42 loc) · 1.9 KB
/
MyTest.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
import torch
import torch.nn.functional as F
import numpy as np
import os
import argparse
from scipy import misc # NOTES: pip install scipy == 1.2.2
from Src.SINet import SINet_ResNet50
from Src.utils.Dataloader import test_dataset
from Src.utils.trainer import eval_mae, numpy2tensor
parser = argparse.ArgumentParser()
parser.add_argument('--testsize', type=int, default=352, help='the snapshot input size')
parser.add_argument('--model_path', type=str,
default='./Snapshot/2020-CVPR-SINet/SINet_40.pth')
parser.add_argument('--test_save', type=str,
default='./Result/2020-CVPR-SINet-New/')
opt = parser.parse_args()
model = SINet_ResNet50().cuda()
model.load_state_dict(torch.load(opt.model_path))
model.eval()
for dataset in ['CAMO', 'CHAMELEON', 'CPD1K', 'COD10K']:
save_path = opt.test_save + dataset + '/'
os.makedirs(save_path, exist_ok=True)
test_loader = test_dataset('./Dataset/TestDataset/{}/Image/'.format(dataset),
'./Dataset/TestDataset/{}/GT/'.format(dataset), opt.testsize)
img_count = 1
for iteration in range(test_loader.size):
image, gt, name = test_loader.load_data()
gt = np.asarray(gt, np.float32)
gt /= (gt.max() + 1e-8)
image = image.cuda()
_, cam = model(image)
cam = F.upsample(cam, size=gt.shape, mode='bilinear', align_corners=True)
cam = cam.sigmoid().data.cpu().numpy().squeeze()
# normalize
cam = (cam - cam.min()) / (cam.max() - cam.min() + 1e-8)
misc.imsave(save_path+name, cam)
mae = eval_mae(numpy2tensor(cam), numpy2tensor(gt))
# coarse score
print('[Eval-Test] Dataset: {}, Image: {} ({}/{}), MAE: {}'.format(dataset, name, img_count,
test_loader.size, mae))
img_count += 1
print("\n[Congratulations! Testing Done]")