forked from dddb11/MVSS-Net
-
Notifications
You must be signed in to change notification settings - Fork 0
/
inference.py
85 lines (73 loc) · 2.97 KB
/
inference.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
import sys
import os
import numpy as np
from common.utils import Progbar, read_annotations
import torch.backends.cudnn as cudnn
from models.mvssnet import get_mvss
from models.resfcn import ResFCN
import torch.utils.data
from common.tools import inference_single
import cv2
#from apex import amp
import argparse
def get_opt():
parser = argparse.ArgumentParser(description='Inference')
parser.add_argument("--model_path", type=str, help="Path to the pretrained model", default="ckpt/mvssnet_github.pt")
parser.add_argument("--test_file", default="./CASIAv1.txt", type=str, help="Path to the image list")
parser.add_argument("--save_dir", type=str, default="./save")
parser.add_argument("--resize", type=int, default=512)
opt = parser.parse_args()
return opt
if __name__ == '__main__':
opt = get_opt()
print("in the head of inference:", opt)
cudnn.benchmark = True
# read test data
test_file = opt.test_file
dataset_name = os.path.basename(test_file).split('.')[0]
model_type = os.path.basename(opt.model_path).split('.')[0]
if not os.path.exists(test_file):
print("%s not exists,quit" % test_file)
sys.exit()
test_data = read_annotations(test_file)
new_size = opt.resize
# load model
model_path = opt.model_path
model = get_mvss(backbone='resnet50',
pretrained_base=True,
nclass=1,
sobel=True,
constrain=True,
n_input=3,
)
if os.path.exists(model_path):
checkpoint = torch.load(model_path, map_location='cpu')
model.load_state_dict(checkpoint, strict=True)
model.eval()
print("load %s finish" % (os.path.basename(model_path)))
else:
print("%s not exist" % model_path)
sys.exit()
model.cuda()
# amp.register_float_function(torch, 'sigmoid')
# model = amp.initialize(models=model, opt_level='O1', loss_scale='dynamic')
model.eval()
save_path = os.path.join(opt.save_dir, dataset_name, model_type)
print("predicted maps will be saved in :%s" % save_path)
os.makedirs(save_path, exist_ok=True)
with torch.no_grad():
progbar = Progbar(len(test_data), stateful_metrics=['path'])
pd_img_lab = []
lab_all = []
scores = []
for ix, (img_path, _, _) in enumerate(test_data):
img = cv2.imread(img_path)
ori_size = img.shape
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = cv2.resize(img, (new_size, new_size))
seg, _ = inference_single(img=img, model=model, th=0)
save_seg_path = os.path.join(save_path, 'pred', os.path.split(img_path)[-1].split('.')[0] + '.png')
os.makedirs(os.path.split(save_seg_path)[0], exist_ok=True)
seg = cv2.resize(seg, (ori_size[1], ori_size[0]))
cv2.imwrite(save_seg_path, seg.astype(np.uint8))
progbar.add(1, values=[('path', save_seg_path), ])