-
Notifications
You must be signed in to change notification settings - Fork 9
/
test.py
71 lines (55 loc) · 2.52 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
#!/usr/bin/python3
#coding=utf-8
import os
os.environ['CUDA_VISIBLE_DEVICES'] = "2"
import sys
sys.path.insert(0, '../')
sys.dont_write_bytecode = True
import cv2
import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn.functional as F
import dataset_1st
from torch.utils.data import DataLoader
from models.modeling import VisionTransformer, get_config
class Test(object):
def __init__(self, Dataset, Network, Path, snapshot):
## dataset
self.cfg = Dataset.Config(datapath=Path, snapshot=snapshot, mode='test')
self.data = Dataset.Data(self.cfg)
self.loader = DataLoader(self.data, batch_size=1, shuffle=False, num_workers=8)
config = get_config()
self.net = Network(config, img_size=352, zero_head=False)
self.net.cuda()
model_dict = self.net.state_dict()
pretrained_dict = torch.load(self.cfg.snapshot,map_location=torch.device('cpu'))
pretrained_dict = {k.replace('module.', ''): v for k, v in pretrained_dict.items() if (k.replace('module.', '') in model_dict)}
# check unloaded weights
for k,v in model_dict.items():
if k in pretrained_dict.keys():
pass
else:
print("miss keys in pretrained_dict: {}".format(k))
model_dict.update(pretrained_dict)
self.net.load_state_dict(model_dict)
self.net.train(False)
def save(self):
with torch.no_grad():
for image, (H, W), name in self.loader:
image, shape = image.cuda().float(), (H, W)
image = F.interpolate(image, (352, 352), mode='bilinear', align_corners=True)
out, _ = self.net(image)
pred = torch.sigmoid(out[0, 0]).cpu().numpy() * 255
pred = cv2.resize(pred, dsize=(W,H), interpolation=cv2.INTER_LINEAR)
head = './eval/maps/' + self.cfg.datapath.split('/')[-1]
if not os.path.exists(head):
os.makedirs(head)
cv2.imwrite(head + '/' + name[0] + '.png', np.round(pred))
if __name__=='__main__':
for path in [ '/home/gaosy/DATA/GT/DUT_O',
'/home/gaosy/DATA/GT/ECSSD', '/home/gaosy/DATA/GT/PASCAL_S',
'/home/gaosy/DATA/GT/DUTS_test', '/home/gaosy/DATA/GT/HKU_IS']: # '/home/gaosy/DATA/DUTS/DUTS-TR'
t = Test(dataset_1st, VisionTransformer, path, './out_2nd/'+'model-x')
# t = Test(dataset_1st, VisionTransformer, path, './out_2nd/'+'model-final')
t.save()