def main(): net = BDRAR().cuda() #if len(args['snapshot']) > 0: # print 'load snapshot \'%s\' for testing' % args['snapshot'] net.load_state_dict(torch.load('model.pth', map_location='cuda:0')) net.eval() with torch.no_grad(): #for name, root in to_test.iteritems(): img_list = ["shadow1.jpg", "shadow2.jpg", "shadow3.jpg"] #img_list = [img_name for img_name in os.listdir(os.path.join(root, 'ShadowImages')) if # img_name.endswith('.jpg')] for idx, img_name in enumerate(img_list): print('predicting for %s: %d / %d' % (img_name, idx + 1, len(img_list))) #check_mkdir( # os.path.join(ckpt_path, exp_name, '(%s) %s_prediction_%s' % (exp_name, name, args['snapshot']))) img = Image.open(img_name) w, h = img.size img_var = Variable(img_transform(img).unsqueeze(0)).cuda() res = net(img_var) prediction = np.array( transforms.Resize((h, w))(to_pil(res.data.squeeze(0).cpu()))) #prediction = crf_refine(np.array(img.convert('RGB')), prediction) Image.fromarray(prediction).save("out_" + img_name)
def main(): net = BDRAR().cuda() if len(args['snapshot']) > 0: print 'load snapshot \'%s\' for testing' % args['snapshot'] net.load_state_dict( torch.load( os.path.join(ckpt_path, exp_name, args['snapshot'] + '.pth'))) test_tp, test_fp, test_fn = 0., 0., 0. net.eval() with torch.no_grad(): for name, root in to_test.iteritems(): img_list = sorted([ img_name for img_name in os.listdir(os.path.join(root, 'images')) if img_name.endswith('.jpg') ]) gt_list = sorted([ img_name for img_name in os.listdir(os.path.join(root, 'masks')) if img_name.endswith('.png') ]) assert len(img_list) == len(gt_list), "%d != %d".format( len(img_list), len(gt_list)) for idx, (img_name, gt_name) in enumerate(zip(img_list, gt_list)): print 'predicting for %s: %d / %d' % (name, idx + 1, len(img_list)) check_mkdir( os.path.join( ckpt_path, exp_name, '(%s) %s_prediction_%s' % (exp_name, name, args['snapshot']))) img = Image.open(os.path.join(root, 'images', img_name)) gt = Image.open(os.path.join(root, 'masks', gt_name)) w, h = img.size img_var = Variable(img_transform(img).unsqueeze(0)).cuda() res = net(img_var) prediction = np.array( transforms.Resize( (h, w))(to_pil(res.data.squeeze(0).cpu()))) prediction = crf_refine(np.array(img.convert('RGB')), prediction) pred_binary = (np.array(prediction) > 0).astype(int) gt = (np.array(gt) > 0).astype(int) tp, fp, fn = compute_stats(pred_binary, gt) test_tp += tp test_fp += fp test_fn += fn test_prec = test_tp / (test_tp + test_fp) test_rec = test_tp / (test_tp + test_fn) test_f1 = 2 * test_prec * test_rec / (test_prec + test_rec) print('Testing percision: %.2f, recall: %.2f, f1: %.2f' % (test_prec, test_rec, test_f1))
def main(): net = BDRAR().cuda() if len(args['snapshot']) > 0: print 'load snapshot \'%s\' for testing' % args['snapshot'] net.load_state_dict( torch.load( os.path.join(ckpt_path, exp_name, args['snapshot'] + '.pth'))) net.eval() with torch.no_grad(): for name, root in to_test.iteritems(): img_list = [ img_name for img_name in os.listdir(os.path.join(root, 'images')) if img_name.endswith('.jpg') ] for idx, img_name in enumerate(img_list): print 'predicting for %s: %d / %d' % (name, idx + 1, len(img_list)) check_mkdir( os.path.join( ckpt_path, exp_name, '(%s) %s_prediction_%s' % (exp_name, name, args['snapshot']))) img = Image.open(os.path.join(root, 'images', img_name)) w, h = img.size img_var = Variable(img_transform(img).unsqueeze(0)).cuda() res = net(img_var) prediction = np.array( transforms.Resize( (h, w))(to_pil(res.data.squeeze(0).cpu()))) prediction = crf_refine(np.array(img.convert('RGB')), prediction) Image.fromarray(prediction).save( os.path.join( ckpt_path, exp_name, '(%s) %s_prediction_%s' % (exp_name, name, args['snapshot']), os.path.splitext(img_name)[0] + ".png"))
def main(): net = BDRAR().cuda().train() optimizer = optim.SGD([{ 'params': [ param for name, param in net.named_parameters() if name[-4:] == 'bias' ], 'lr': 2 * args['lr'] }, { 'params': [ param for name, param in net.named_parameters() if name[-4:] != 'bias' ], 'lr': args['lr'], 'weight_decay': args['weight_decay'] }], momentum=args['momentum']) if len(args['snapshot']) > 0: print('training resumes from \'%s\'' % args['snapshot']) net.load_state_dict( torch.load( os.path.join(ckpt_path, exp_name, args['snapshot'] + '.pth'))) optimizer.load_state_dict( torch.load( os.path.join(ckpt_path, exp_name, args['snapshot'] + '_optim.pth'))) optimizer.param_groups[0]['lr'] = 2 * args['lr'] optimizer.param_groups[1]['lr'] = args['lr'] check_mkdir(ckpt_path) check_mkdir(os.path.join(ckpt_path, exp_name)) open(log_path, 'w').write(str(args) + '\n\n') train(net, optimizer)