-
Notifications
You must be signed in to change notification settings - Fork 2
/
Copy pathtest.py
68 lines (49 loc) · 1.56 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
import torch
import torch.nn.functional as F
import sys
sys.path.append('./models')
import numpy as np
import os, argparse
import cv2
from net import SwinMCNet
from data import test_dataset
from options import opt
from collections import OrderedDict
import time
from os.path import splitext
from ptflops.flops_counter import get_model_complexity_info
#set device for test
os.environ["CUDA_VISIBLE_DEVICES"] = "3"
#load the model
model = SwinMCNet()
base_weights = torch.load(opt.test_model)
new_state_dict = OrderedDict()
for k, v in base_weights.items():
name = k[7:] # remove 'module.'
new_state_dict[name] = v
model.load_state_dict(new_state_dict)
print('Loading base network...')
model.cuda()
model.eval()
#test
test_data_root = opt.test_data_root
maps_path = opt.maps_path
test_sets = ['VT5000/Test','VT1000','VT821']
for dataset in test_sets:
save_path = maps_path + dataset + '/'
if not os.path.exists(save_path):
os.makedirs(save_path)
dataset_path = test_data_root + dataset
test_loader = test_dataset(dataset_path, opt.testsize)
for i in range(test_loader.size):
image, t, gt, (H, W), name = test_loader.load_data()
image = image.cuda()
t = t.cuda()
shape = (W,H)
outi1, outt1, out1, outi2, outt2, out2 = model(image,t,shape)
res = out2
res = res.sigmoid().data.cpu().numpy().squeeze()
res = (res - res.min()) / (res.max() - res.min() + 1e-8)
print('save img to: ',save_path + name)
cv2.imwrite(save_path + name,res*255)
print('Test Done!')