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))