-
Notifications
You must be signed in to change notification settings - Fork 1
/
mynet_test.py
83 lines (62 loc) · 2.09 KB
/
mynet_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
from skimage import io
import torch
from torch.autograd import Variable
from torch.utils.data import DataLoader
from torchvision import transforms
from PIL import Image
import glob
from tqdm import tqdm
from data_loader import Rescale
from data_loader import ToTensor
from data_loader import SalObjDataset
from model import FSMINet
def normPRED(x):
MAX = torch.max(x)
MIN = torch.min(x)
out = (x-MIN)/(MAX-MIN)
return out
def save_output(image_name,pred,d_dir):
predict = pred
predict = predict.squeeze()
predict_np = predict.cpu().data.numpy()
im = Image.fromarray(predict_np*255).convert('RGB')
img_name = image_name.split("/")[-1]
image = io.imread(image_name)
imo = im.resize((image.shape[1],image.shape[0]),resample=Image.BILINEAR)
t = img_name.split(".")
t = t[0:-1]
imidx = t[0]
for i in range(1,len(t)):
imidx = imidx + "." + t[i]
imo.save(d_dir+imidx+'.png')
if __name__ == '__main__':
# --------- Define the address and image format ---------
image_dir = ""
prediction_dir = ""
model_dir = ""
img_name_list = glob.glob(image_dir + '*.jpg')
# --------- Load the data ---------
test_salobj_dataset = SalObjDataset(img_name_list = img_name_list, lbl_name_list = [],transform=transforms.Compose([Rescale(384),ToTensor(flag=0)]))
test_salobj_dataloader = DataLoader(test_salobj_dataset, batch_size=1,shuffle=False,num_workers=1)
# --------- Define the model ---------
print("...load FSMINet...")
net = FSMINet()
net.load_state_dict(torch.load(model_dir))
if torch.cuda.is_available():
net.cuda()
net.eval()
# --------- Generate prediction images ---------
for i_test, data_test in tqdm(enumerate(test_salobj_dataloader)):
inputs_test = data_test['image']
inputs_test = inputs_test.type(torch.FloatTensor)
if torch.cuda.is_available():
inputs_test = Variable(inputs_test.cuda())
else:
inputs_test = Variable(inputs_test)
d0, d1, d2, d3, d4, d5 = net(inputs_test)
# normalization
pred = d0[:,0,:,:]
pred = normPRED(pred)
# save results to test_results folder
save_output(img_name_list[i_test],pred,prediction_dir)
del d0, d1, d2, d3, d4, d5