def NLFD(model_path, img_path, cuda):
transform = transforms.Compose(
[transforms.Resize((352, 352)),
transforms.ToTensor()])
img = Image.open(img_path)
shape = img.size
ori_H = shape[0]
ori_W = shape[1]
img = transform(img) - torch.Tensor([123.68, 116.779, 103.939]).view(
3, 1, 1) / 255
img = Variable(img.unsqueeze(0), volatile=True)
net = build_model()
net.load_state_dict(torch.load(model_path))
net.eval()
if cuda: img, net = img.cuda(), net.cuda()
prob = net(img)
prob = (prob.cpu().data[0][0].numpy() * 255).astype(np.uint8)
# kernel = np.ones((5, 5), np.uint8)
# prob = cv2.dilate(prob, kernel, iterations=1)
#print (prob)
#prob = cv2.morphologyEx(prob, cv2.MORPH_CLOSE, kernel)
# p_img = Image.fromarray(prob, mode='L').resize(shape)
# p_img.show()
prob = resize(prob, (ori_W, ori_H), order=1)
#prob = cv2.cvtColor(prob.copy(), cv2.COLOR_GRAY2BGR)
return prob
def build_model(self):
self.net = build_model()
if self.config.mode == 'train':
self.loss = Loss(self.config.area, self.config.boundary)
if self.config.cuda: self.net = self.net.cuda()
if self.config.cuda and self.config.mode == 'train':
self.loss = self.loss.cuda()
self.net.train()
self.net.apply(weights_init)
if self.config.load == '':
self.net.base.load_state_dict(torch.load(self.config.vgg))
if self.config.load != '':
self.net.load_state_dict(torch.load(self.config.load))
self.optimizer = Adam(self.net.parameters(), self.config.lr)
self.print_network(self.net, 'NLFD')
def demo(model_path, img_path, cuda):
transform = transforms.Compose(
[transforms.Resize((352, 352)),
transforms.ToTensor()])
img = Image.open(img_path)
shape = img.size
img = transform(img) - torch.Tensor([123.68, 116.779, 103.939]).view(
3, 1, 1) / 255
img = Variable(img.unsqueeze(0), volatile=True)
net = build_model()
net.load_state_dict(torch.load(model_path))
net.eval()
if cuda: img, net = img.cuda(), net.cuda()
prob = net(img)
prob = (prob.cpu().data[0][0].numpy() * 255).astype(np.uint8)
p_img = Image.fromarray(prob, mode='L').resize(shape)
p_img.show()
img = img.cuda()
prob = net(img)
prob = prob.cpu().data[0][0].numpy()
prob = postprocess(prob)
prob = cv2.resize(prob, (img_w, img_h)).astype(np.uint8)
return prob
if __name__ == '__main__':
model_path = sys.argv[1]
img_path = sys.argv[2]
cuda = False
img = cv2.imread(img_path)
net = build_model()
net.load_state_dict(torch.load(model_path))
net.eval()
if cuda:
net = net.cuda()
num_run = 1
total_time = 0
for i in range(num_run):
start = time.time()
mask = demo(net, img, cuda)
eps = time.time() - start
print('run-%d\t%.3f' % (i, eps * 1000))
if i != 0: # exclude first run
total_time += eps
#print('average runtime[%s] %.3fms' % (mode, total_time/(num_run-1)*1000))