def main():
net = DAF().cuda()
if len(args['snapshot']) > 0:
print('training resumes from ' + args['snapshot'])
net.load_state_dict(
torch.load(
os.path.join(ckpt_path, exp_name, args['snapshot'] + '.pth')))
net.eval()
for idx, img_name in enumerate(os.listdir(os.path.join(testing_root,
'us'))):
print 'predicting %d' % (idx + 1)
check_mkdir(
os.path.join(ckpt_path, exp_name,
'prediction_' + args['snapshot']))
img = Image.open(os.path.join(testing_root, 'us',
img_name)).convert('RGB')
img_var = Variable(img_transform(img).unsqueeze(0)).cuda()
prediction = np.array(to_pil(net(img_var).data.squeeze(0).cpu()))
prediction = crf_refine(np.array(img), prediction)
Image.fromarray(prediction).save(
os.path.join(ckpt_path, exp_name, 'prediction_' + args['snapshot'],
img_name))
def run_test(ckp_name):
sess = Session()
sess.net.eval()
sess.load_checkpoints(ckp_name,'test')
if sess.multi_gpu :
sess.net = nn.DataParallel(sess.net)
sess.batch_size = 1
sess.shuffle = False
sess.outs = -1
dt = sess.get_dataloader(sess.test_data_path, train_mode=False)
input_names = open(sess.test_data_path+'SBU.txt').readlines()
widgets = [progressbar.Percentage(),progressbar.Bar(),progressbar.ETA()]
bar = progressbar.ProgressBar(widgets=widgets,maxval=len(dt)).start()
for i, batch in enumerate(dt):
pred = sess.inf_batch('test', batch)
image = I.open(sess.test_data_path+input_names[i].split(' ')[0]).convert('RGB')
final = I.fromarray((pred[-1].cpu().data * 255).numpy().astype('uint8')[0,0,:,:])
final = np.array(final.resize(image.size))
final_crf = crf_refine(np.array(image),final)
ensure_dir('./results')
io.imsave('./results/'+input_names[i].split(' ')[0].split('/')[1][:-3]+'png',final_crf)
bar.update(i+1)
def main():
net = DSDNet()#.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.items():
img_list = [img_name for img_name in os.listdir(root) 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('Output/SBU', '%s_%s_prediction_%s' % (exp_name, name, args['snapshot'])))
img = Image.open(os.path.join(root, img_name))
w, h = img.size
img_var = Variable(img_transform(img).unsqueeze(0)).cpu()
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('Output/SBU', '%s_%s_prediction_%s' % (
exp_name, name, args['snapshot']), img_name[:-4])+'.png')
def main():
net = TAYLOR5().cuda(device_ids[0])
if len(args['snapshot']) > 0:
print('Load snapshot {} for testing'.format(args['snapshot']))
net.load_state_dict(torch.load(os.path.join(ckpt_path, exp_name, args['snapshot'] + '.pth')))
print('Load {} succeed!'.format(os.path.join(ckpt_path, exp_name, args['snapshot'] + '.pth')))
net.eval()
with torch.no_grad():
for name, root in to_test.items():
# img_list = [img_name for img_name in os.listdir(os.path.join(root, 'image')) if img_name.endswith('.jpg')]
img_list = [img_name for img_name in os.listdir(os.path.join(root, 'image'))]
start = time.time()
for idx, img_name in enumerate(img_list):
print('predicting for {}: {:>4d} / {}'.format(name, idx + 1, len(img_list)))
check_mkdir(os.path.join(ckpt_path, exp_name, '%s_%s' % (exp_name, args['snapshot'])))
img = Image.open(os.path.join(root, 'image', img_name))
if img.mode != 'RGB':
img = img.convert('RGB')
print("{} is a gray image.".format(name))
w, h = img.size
img_var = Variable(img_transform(img).unsqueeze(0)).cuda()
# f_4, f_3, f_2, f_1, e = net(img_var)
f_4, f_3, f_2, f_1 = net(img_var)
# output = f.data.squeeze(0).cpu()
# edge = e.data.squeeze(0).cpu()
f_4 = f_4.data.squeeze(0).cpu()
f_3 = f_3.data.squeeze(0).cpu()
f_2 = f_2.data.squeeze(0).cpu()
f_1 = f_1.data.squeeze(0).cpu()
# prediction = np.array(transforms.Resize((h, w))(to_pil(output)))
# edge = np.array(transforms.Resize((h, w))(to_pil(edge)))
f_4 = np.array(transforms.Resize((h, w))(to_pil(f_4)))
f_3 = np.array(transforms.Resize((h, w))(to_pil(f_3)))
f_2 = np.array(transforms.Resize((h, w))(to_pil(f_2)))
f_1 = np.array(transforms.Resize((h, w))(to_pil(f_1)))
if args['crf']:
# prediction = crf_refine(np.array(img.convert('RGB')), prediction)
f_1 = crf_refine(np.array(img.convert('RGB')), f_1)
# f_2 = crf_refine(np.array(img.convert('RGB')), f_2)
# f_3 = crf_refine(np.array(img.convert('RGB')), f_3)
# f_4 = crf_refine(np.array(img.convert('RGB')), f_4)
# edge = crf_refine(np.array(img.convert('RGB')), edge)
# Image.fromarray(prediction).save(os.path.join(ckpt_path, exp_name, '%s_%s' % (exp_name, args['snapshot']), 'map', img_name[:-4] + ".png"))
# Image.fromarray(edge).save(os.path.join(ckpt_path, exp_name, '%s_%s' % (exp_name, args['snapshot']), 'edge', img_name[:-4] + ".png"))
# Image.fromarray(f_4).save(os.path.join(ckpt_path, exp_name, '%s_%s' % (exp_name, args['snapshot']), 'f4', img_name[:-4] + ".png"))
# Image.fromarray(f_3).save(os.path.join(ckpt_path, exp_name, '%s_%s' % (exp_name, args['snapshot']), 'f3', img_name[:-4] + ".png"))
# Image.fromarray(f_2).save(os.path.join(ckpt_path, exp_name, '%s_%s' % (exp_name, args['snapshot']), 'f2', img_name[:-4] + ".png"))
Image.fromarray(f_1).save(os.path.join(ckpt_path, exp_name, '%s_%s' % (exp_name, args['snapshot']), 'f1', img_name[:-4] + ".png"))
# skimage.io.imsave(os.path.join(ckpt_path, exp_name, '%s_%s' % (exp_name, args['snapshot']), 'f1', img_name[:-4] + ".png"), np.where(f_1>=127.5, 255, 0).astype(np.uint8))
# skimage.io.imsave(os.path.join(ckpt_path, exp_name, '%s_%s' % (exp_name, args['snapshot']), img_name[:-4] + ".png"), prediction.astype(np.uint8))
# check_mkdir(os.path.join(msd_testing_root, 'taylor5_' + str(args['scale'])))
# Image.fromarray(f_1).save(os.path.join(msd_testing_root, 'taylor5_' + str(args['scale']), img_name[:-4] + ".png"))
end = time.time()
print("Average Time Is : {:.2f}".format((end - start) / len(img_list)))
def main():
net = R3Net().cuda()
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()
results = {}
with torch.no_grad():
for name, root in to_test.iteritems():
precision_record, recall_record, = [
AvgMeter() for _ in range(256)
], [AvgMeter() for _ in range(256)]
mae_record = AvgMeter()
if args['save_results']:
check_mkdir(
os.path.join(
ckpt_path, exp_name,
'(%s) %s_%s' % (exp_name, name, args['snapshot'])))
img_list = os.listdir(root)
print img_list
for idx, img_name in enumerate(img_list):
print 'predicting for %s: %d / %d' % (name, idx + 1,
len(img_list))
#print img_name
img_path = os.path.join(root, img_name)
img = Image.open(img_path).convert('RGB')
img_var = Variable(img_transform(img).unsqueeze(0),
volatile=True).cuda()
prediction = net(img_var)
prediction = np.array(to_pil(prediction.data.squeeze(0).cpu()))
if args['crf_refine']:
prediction = crf_refine(np.array(img), prediction)
#gt = np.array(Image.open(os.path.join(root+"/masks", img_name )).convert('L'))
#precision, recall, mae = cal_precision_recall_mae(prediction, gt)
#for pidx, pdata in enumerate(zip(precision, recall)):
# p, r = pdata
# precision_record[pidx].update(p)
# recall_record[pidx].update(r)
#mae_record.update(mae)
if args['save_results']:
Image.fromarray(prediction).save(
os.path.join(
ckpt_path, exp_name, '%s_%s_%s' %
(exp_name, name, args['snapshot'] + "kaist_test"),
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():
#########################Load##########################
net = GateNet().cuda()
# net = GateNet_SIM_Light().cuda() # vgg16
print('load snapshot \'%s\' for testing' % args['snapshot'])
net.load_state_dict(
torch.load(os.path.join(ckpt_path, exp_name, args['snapshot']),
map_location={'cuda:1': 'cuda:1'}))
net.eval()
# transforms = tta.Compose(
# [
# tta.HorizontalFlip(),
# tta.Scale(scales=[1,1.5],interpolation='bilinear',align_corners=False),
# ]
# )
#
# net = tta.SegmentationTTAWrapper(net, transforms, merge_mode='mean')
########################################################
with torch.no_grad():
for name, root in to_test.items():
check_mkdir(
os.path.join(ckpt_path, exp_name, '(%s) %s_%s' %
(exp_name, name, args['snapshot'])))
root1 = os.path.join(root, 'GT')
img_list = [
os.path.splitext(f)[0] for f in os.listdir(root1)
if f.endswith('.png')
]
print(img_list)
for idx, img_name in enumerate(img_list):
print('predicting for %s: %d / %d' %
(name, idx + 1, len(img_list)))
img1 = Image.open(
os.path.join(root,
'Imgs/' + img_name + '.jpg')).convert('RGB')
img = img1
w, h = img1.size
img1 = img1.resize([384, 384], Image.BILINEAR)
img_var = Variable(img_transform(img1).unsqueeze(0),
volatile=True).cuda()
prediction = net(img_var)
# prediction = F.sigmoid(prediction)
prediction = to_pil(prediction.data.squeeze(0).cpu())
# prediction = prediction.resize((w, h), Image.BILINEAR)
prediction = prediction.resize((w, h), Image.NEAREST)
if args['crf_refine']:
prediction = crf_refine(np.array(img),
np.array(prediction))
prediction = np.array(prediction)
if args['save_results']:
Image.fromarray(prediction).save(
os.path.join(ckpt_path, exp_name, 'DUTS',
img_name + '.png'))
def main():
t0 = time.time()
net = RGBD_sal().cuda()
print('load snapshot \'%s\' for testing' % args['snapshot'])
net.load_state_dict(
torch.load(os.path.join(ckpt_path, exp_name,
args['snapshot'] + '.pth'),
map_location={'cuda:1': 'cuda:1'}))
net.eval()
with torch.no_grad():
for name, root in to_test.items():
check_mkdir(
os.path.join(ckpt_path, exp_name, '(%s) %s_%s' %
(exp_name, name, args['snapshot'])))
root1 = os.path.join(root, 'GT')
# root2 = os.path.join('/home/zxq/桌面/RGBD_saliencydata_results/2/model_fpn_rgbd_depth_enhanced_attention_saliency_background_branch_vgg16/NJUD')
img_list = [
os.path.splitext(f)[0] for f in os.listdir(root1)
if f.endswith('.png')
]
for idx, img_name in enumerate(img_list):
print('predicting for %s: %d / %d' %
(name, idx + 1, len(img_list)))
img1 = Image.open(os.path.join(root, 'RGB', img_name +
'.jpg')).convert('RGB')
depth = Image.open(
os.path.join(root, 'depth',
img_name + '.png')).convert('L')
img = img1
w_, h_ = img1.size
img1 = img1.resize([384, 384])
depth = depth.resize([384, 384])
img_var = Variable(img_transform(img1).unsqueeze(0),
volatile=True).cuda()
depth = Variable(depth_transform(depth).unsqueeze(0),
volatile=True).cuda()
prediction = net(img_var, depth)
prediction = to_pil(prediction.data.squeeze(0).cpu())
prediction = prediction.resize((w_, h_), Image.BILINEAR)
if args['crf_refine']:
prediction = crf_refine(np.array(img),
np.array(prediction))
prediction = np.array(prediction)
if args['save_results']:
Image.fromarray(prediction).save(
os.path.join(
ckpt_path, exp_name, '(%s) %s_%s' % (
exp_name,
name,
args['snapshot'],
), img_name + '.png'))
def main():
net = MirrorNet().cuda(device_ids[0])
if len(args['snapshot']) > 0:
print('Load snapshot {} for testing'.format(args['snapshot']))
net.load_state_dict(
torch.load(os.path.join(ckpt_path, exp_name, 'MirrorNet.pth')))
print('Load {} succeed!'.format(
os.path.join(ckpt_path, exp_name, 'MirrorNet.pth')))
net.eval()
with torch.no_grad():
for name, root in to_test.items():
img_list = [
img_name
for img_name in os.listdir(os.path.join(root, 'image'))
]
start = time.time()
for idx, img_name in enumerate(img_list):
print('predicting for {}: {:>4d} / {}'.format(
name, idx + 1, len(img_list)))
check_mkdir(
os.path.join(ckpt_path, exp_name,
'%s_%s' % (exp_name, args['snapshot'])))
img = Image.open(os.path.join(root, 'image', img_name))
if img.mode != 'RGB':
img = img.convert('RGB')
print("{} is a gray image.".format(name))
w, h = img.size
img_var = Variable(img_transform(img).unsqueeze(0)).cuda(
device_ids[0])
f_4, f_3, f_2, f_1 = net(img_var)
f_4 = f_4.data.squeeze(0).cpu()
f_3 = f_3.data.squeeze(0).cpu()
f_2 = f_2.data.squeeze(0).cpu()
f_1 = f_1.data.squeeze(0).cpu()
f_4 = np.array(transforms.Resize((h, w))(to_pil(f_4)))
f_3 = np.array(transforms.Resize((h, w))(to_pil(f_3)))
f_2 = np.array(transforms.Resize((h, w))(to_pil(f_2)))
f_1 = np.array(transforms.Resize((h, w))(to_pil(f_1)))
if args['crf']:
f_1 = crf_refine(np.array(img.convert('RGB')), f_1)
Image.fromarray(f_1).save(
os.path.join(ckpt_path, exp_name,
'%s_%s' % (exp_name, args['snapshot']),
img_name[:-4] + ".png"))
end = time.time()
print("Average Time Is : {:.2f}".format(
(end - start) / len(img_list)))
def main():
net = EDGE_CBAM_X_CCL().cuda()
if len(args['snapshot']) > 0:
print('Load snapshot {} for testing'.format(args['snapshot']))
net.load_state_dict(
torch.load(
os.path.join(ckpt_path, exp_name, args['snapshot'] + '.pth')))
print('Load {} succeed!'.format(
os.path.join(ckpt_path, exp_name, args['snapshot'] + '.pth')))
net.eval()
with torch.no_grad():
for name, root in to_test.items():
img_list = [
img_name
for img_name in os.listdir(os.path.join(root, 'image'))
if img_name.endswith('.jpg')
]
start = time.time()
for idx, img_name in enumerate(img_list):
print('predicting for {}: {:>4d} / {}'.format(
name, idx + 1, len(img_list)))
check_mkdir(
os.path.join(ckpt_path, exp_name,
'%s_%s' % (exp_name, args['snapshot'])))
img = Image.open(os.path.join(root, 'image', img_name))
w, h = img.size
img_var = Variable(img_transform(img).unsqueeze(0)).cuda()
edge, res = net(img_var)
edge = np.array(
transforms.Resize(
(h, w))(to_pil(edge.data.squeeze(0).cpu())))
prediction = np.array(
transforms.Resize(
(h, w))(to_pil(res.data.squeeze(0).cpu())))
if args['crf']:
prediction = crf_refine(np.array(img.convert('RGB')),
prediction)
Image.fromarray(edge).save(
os.path.join(ckpt_path, exp_name,
'%s_%s' % (exp_name, args['snapshot']),
img_name[:-4] + "_edge.png"))
Image.fromarray(prediction).save(
os.path.join(ckpt_path, exp_name,
'%s_%s' % (exp_name, args['snapshot']),
img_name[:-4] + ".png"))
end = time.time()
print("Average Time Is : {:.2f}".format(
(end - start) / len(img_list)))
def main():
net = PMD().cuda(device_ids[0])
net.load_state_dict(torch.load('pmd.pth'))
net.eval()
with torch.no_grad():
for name, root in to_test.items():
img_list = [
img_name
for img_name in os.listdir(os.path.join(root, 'image'))
]
start = time.time()
for idx, img_name in enumerate(img_list):
print('predicting for {}: {:>4d} / {}'.format(
name, idx + 1, len(img_list)))
check_mkdir(os.path.join("result", name))
img = Image.open(os.path.join(root, 'image', img_name))
if img.mode != 'RGB':
img = img.convert('RGB')
print("{} is a gray image.".format(name))
w, h = img.size
img_var = Variable(img_transform(img).unsqueeze(0)).cuda()
# f_4, f_3, f_2, f_1, e = net(img_var)
f_4, f_3, f_2, f_1, edge, final = net(img_var)
# output = f.data.squeeze(0).cpu()
# edge = e.data.squeeze(0).cpu()
f_4 = f_4.data.squeeze(0).cpu()
f_3 = f_3.data.squeeze(0).cpu()
f_2 = f_2.data.squeeze(0).cpu()
f_1 = f_1.data.squeeze(0).cpu()
edge = edge.data.squeeze(0).cpu()
final = final.data.squeeze(0).cpu()
f_4 = np.array(transforms.Resize((h, w))(to_pil(f_4)))
f_3 = np.array(transforms.Resize((h, w))(to_pil(f_3)))
f_2 = np.array(transforms.Resize((h, w))(to_pil(f_2)))
f_1 = np.array(transforms.Resize((h, w))(to_pil(f_1)))
edge = np.array(transforms.Resize((h, w))(to_pil(edge)))
final = np.array(transforms.Resize((h, w))(to_pil(final)))
if args['crf']:
final = crf_refine(np.array(img.convert('RGB')), final)
Image.fromarray(final).save(
os.path.join("result", name, img_name[:-4] + ".png"))
end = time.time()
print("Average Time Is : {:.2f}".format(
(end - start) / len(img_list)))
def test_step(self, batch, batch_idx):
global test_iter
inputs = batch[0]
outputs = batch[1]
# inputs = torch.from_numpy(inputs)
# outputs = torch.tensor(outputs)
f_4_gpu, f_3_gpu, f_2_gpu, f_1_gpu = self(inputs)
loss1 = lovasz_hinge(f_1_gpu, outputs,
per_image=False) * self.args.w_losses[0]
loss2 = lovasz_hinge(f_2_gpu, outputs,
per_image=False) * self.args.w_losses[1]
loss3 = lovasz_hinge(f_3_gpu, outputs,
per_image=False) * self.args.w_losses[2]
loss4 = lovasz_hinge(f_4_gpu, outputs,
per_image=False) * self.args.w_losses[3]
loss = loss1 + loss2 + loss3 + loss4
self.log('test_loss', loss)
if self.args.developer_mode:
real_img = batch[2]
real_mask = batch[3]
sq_zero = real_img[0].squeeze()
sq_zero = sq_zero.cpu().numpy()
real_img = Image.fromarray(sq_zero)
sq_m = real_mask[0].squeeze()
sq_m = sq_m.cpu().numpy()
real_mask = Image.fromarray(sq_m)
im_size = real_img.size
rev_size = [im_size[1], im_size[0]]
f_1 = f_1_gpu.data.cpu()
f_1_trans = np.array(transforms.Resize(rev_size)(to_pil(f_1[0])))
f_1_crf = crf_refine(np.array(real_img), f_1_trans)
new_image = Image.new('RGB', (3 * im_size[0], im_size[1]),
(250, 250, 250))
img_res = Image.fromarray(f_1_crf)
new_image.paste(real_img, (0, 0))
new_image.paste(real_mask, (im_size[0], 0))
new_image.paste(img_res, (im_size[0] * 2, 0))
# The number of test itteration
# self.test_iter +=1
test_iter += 1
new_image.save(
os.path.join(self.args.msd_results_root, "Testing",
"image: " + str(test_iter) + " test.png"))
# self.logger.experiment.log('val_loss', loss)
return {'test_loss': loss}
def main():
net = OUR2().cuda(device_ids[0])
if len(args['snapshot']) > 0:
print('Load snapshot {} for testing'.format(args['snapshot']))
net.load_state_dict(
torch.load(
os.path.join(ckpt_path, exp_name, args['snapshot'] + '.pth')))
print('Load {} succeed!'.format(
os.path.join(ckpt_path, exp_name, args['snapshot'] + '.pth')))
net.eval()
with torch.no_grad():
for name, root in to_test.items():
# img_list = [img_name for img_name in os.listdir(os.path.join(root, 'image')) if img_name.endswith('.jpg')]
img_list = [
img_name
for img_name in os.listdir(os.path.join(root, 'image'))
]
start = time.time()
for idx, img_name in enumerate(img_list):
print('predicting for {}: {:>4d} / {}'.format(
name, idx + 1, len(img_list)))
check_mkdir(
os.path.join(ckpt_path, exp_name,
'%s_%s' % (exp_name, args['snapshot'])))
img = Image.open(os.path.join(root, 'image', img_name))
if img.mode != 'RGB':
img = img.convert('RGB')
print("{} is a gray image.".format(name))
w, h = img.size
img_var = Variable(img_transform(img).unsqueeze(0)).cuda()
_, _, _, _, f = net(img_var)
output = f.data.squeeze(0).cpu()
prediction = np.array(
transforms.Resize((h, w))(to_pil(output)))
if args['crf']:
prediction = crf_refine(np.array(img.convert('RGB')),
prediction)
# Image.fromarray(prediction).save(os.path.join(ckpt_path, exp_name, '%s_%s' % (exp_name, args['snapshot']), img_name[:-4] + ".png"))
# skimage.io.imsave(os.path.join(ckpt_path, exp_name, '%s_%s' % (exp_name, args['snapshot']), img_name[:-4] + ".png"), np.where(prediction>=127.5, 255, 0).astype(np.uint8))
# skimage.io.imsave(os.path.join(ckpt_path, exp_name, '%s_%s' % (exp_name, args['snapshot']), img_name[:-4] + ".png"), prediction.astype(np.uint8))
check_mkdir(os.path.join(msd_testing_root, 'mirror_map'))
Image.fromarray(prediction).save(
os.path.join(msd_testing_root, 'mirror_map',
img_name[:-4] + ".png"))
end = time.time()
print("Average Time Is : {:.2f}".format(
(end - start) / len(img_list)))
def main():
net = R3Net().cuda()
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()
results = {}
for name, root in to_test.iteritems():
precision_record, recall_record, = [AvgMeter() for _ in range(256)], [AvgMeter() for _ in range(256)]
mae_record = AvgMeter()
img_list = [os.path.splitext(f)[0] for f in os.listdir(root) if f.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_%s' % (exp_name, name, args['snapshot'])))
img = Image.open(os.path.join(root, img_name + '.jpg')).convert('RGB')
img_var = Variable(img_transform(img).unsqueeze(0), volatile=True).cuda()
prediction = net(img_var)
prediction = np.array(to_pil(prediction.data.squeeze(0).cpu()))
if args['crf_refine']:
prediction = crf_refine(np.array(img), prediction)
gt = np.array(Image.open(os.path.join(root, img_name + '.png')).convert('L'))
precision, recall, mae = cal_precision_recall_mae(prediction, gt)
for pidx, pdata in enumerate(zip(precision, recall)):
p, r = pdata
precision_record[pidx].update(p)
recall_record[pidx].update(r)
mae_record.update(mae)
if args['save_results']:
Image.fromarray(prediction).save(os.path.join(ckpt_path, exp_name, '(%s) %s_%s' % (
exp_name, name, args['snapshot']), img_name + '.png'))
fmeasure = cal_fmeasure([precord.avg for precord in precision_record],
[rrecord.avg for rrecord in recall_record])
results[name] = {'fmeasure': fmeasure, 'mae': mae_record.avg}
print
'test results:'
print
results
def test(epoch):
global best_eval
net.eval()
idx = 0
# evaluator = Evaluator_Miou(2)
evaluator_dice = Evaluator_dice()
evaluator_F1 = Evaluator_F1()
with torch.no_grad():
for img_name, gt_name in zip(os.listdir(os.path.join(testing_root, 'us')),
os.listdir(os.path.join(testing_root, 'seg'))):
idx += 1
img = Image.open(os.path.join(testing_root, 'us', img_name)).convert('RGB')
gt = Image.open(os.path.join(testing_root, 'seg', gt_name)).convert('1')
img_var = Variable(transform(img).unsqueeze(0)).cuda()
prediction = np.array(to_pil(net(img_var).data.squeeze(0).cpu()).resize(gt.size))
prediction = crf_refine(np.array(img), prediction)
# a=prediction.sum(axis=0).sum(axis=0)
if flag.save_pred:
check_mkdir(os.path.join(ckpt_path, exp_name, 'prediction'))
Image.fromarray(prediction).save(
os.path.join(ckpt_path, exp_name, 'prediction', img_name))
gt = np.array(gt)
pred = Image.fromarray(prediction).convert('1')
pred = np.array(pred)
evaluator_dice.add_batch(gt, pred)
evaluator_F1.add_batch(pred, gt)
current_dice = evaluator_dice.get_dice()
progress_bar(idx, len(os.listdir(os.path.join(testing_root, 'us'))), 'Dice: %.4f' % (current_dice))
# Miou = evaluator.Frequency_Weighted_Intersection_over_Union()
dice = evaluator_dice.get_dice()
F1 = evaluator_F1.get_F1()
print('Mean dice is %.4f | Mean F1 is %.4f'%(dice, F1))
# Save checkpoint.
if dice > best_eval and not flag.test:
print('Saving..')
state = {
'net': net.state_dict(),
'eval': dice,
'epoch': epoch,
'eval_type': 'dice'
}
checkpoint_path = os.path.join(ckpt_path, exp_name)
if not os.path.isdir(checkpoint_path):
os.mkdir(checkpoint_path)
torch.save(state, os.path.join(checkpoint_path, 'model.pth'))
best_eval = dice
def test(net):
print("Testing Mode ......")
global best_ber
net.eval()
BER = []
for name, root in to_test.items():
img_list = [
img_name for img_name in os.listdir(os.path.join(root, 'image'))
]
for idx, img_name in enumerate(img_list):
# print('predicting for {}: {:>4d} / {}'.format(name, idx + 1, len(img_list)))
# check_mkdir(os.path.join(ckpt_path, exp_name, '%s_%s_%s' % (exp_name, args['snapshot'], 'nocrf')))
img = Image.open(os.path.join(root, 'image', img_name))
gt = Image.open(os.path.join(root, 'mask', img_name[:-4] + '.png'))
gt = np.array(gt)
if img.mode != 'RGB':
img = img.convert('RGB')
print("{} is a gray image.".format(name))
w, h = img.size
img_var = Variable(test_transform(img).unsqueeze(0)).cuda(
device_ids[0])
f_4, f_3, f_2, f_1 = net(img_var)
f_4 = f_4.data.squeeze(0).cpu()
f_3 = f_3.data.squeeze(0).cpu()
f_2 = f_2.data.squeeze(0).cpu()
f_1 = f_1.data.squeeze(0).cpu()
f_4 = np.array(transforms.Resize((h, w))(to_pil(f_4)))
f_3 = np.array(transforms.Resize((h, w))(to_pil(f_3)))
f_2 = np.array(transforms.Resize((h, w))(to_pil(f_2)))
f_1 = np.array(transforms.Resize((h, w))(to_pil(f_1)))
if args['crf']:
f_1 = crf_refine(np.array(img.convert('RGB')), f_1)
gt = np.where(gt == 255, 1, 0).astype(np.float32)
f_1 = np.where(f_1 * 255.0 >= 127.5, 1, 0).astype(np.float32)
ber = compute_ber(f_1, gt)
# print("The %d pics ber is %.5f" % (idx + 1, ber))
BER.append(ber)
mean_BER = 100 * sum(BER) / len(BER)
return mean_BER
def dir_crf_smoothing(self):
imlist = []
imnamelist = []
for root, _, fnames in sorted(os.walk(self.input)):
for fname in fnames:
if fname.endswith(
'.png'
) and 'M1BS' in fname and not fname.startswith('.'):
path = os.path.join(root, fname)
imlist.append((path, fname))
imnamelist.append(fname)
print(imnamelist)
for path, name in imlist:
print(path)
im = misc.imread(path)
raw_pred = misc.imread(os.path.join(self.out_raw, name))
#raw_pred[raw_pred<130] = 0
#raw_pred[raw_pred>=120] = 255
crf_out = crf_refine(im.astype(np.uint8),
raw_pred.astype(np.uint8))
misc.imsave(self.output + '/crf/' + name, crf_out)
def validation_epoch_end(self, outputs):
# outputs = list of dictionaries
avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
if self.args.developer_mode:
batch = self.eval_set.sample(1)
inputs = batch["img"]
outputs = batch["mask"]
inputs = torch.from_numpy(inputs)
outputs = torch.tensor(outputs)
if len(self.args.device_ids) > 0:
inputs = inputs.cuda(self.args.device_ids[0])
outputs = outputs.cuda(self.args.device_ids[0])
_, _, _, f_1_gpu = self(inputs)
f_1 = f_1_gpu.data.cpu()
rev_size = [batch["size"][0][1], batch["size"][0][0]]
image1_size = batch["size"][0]
f_1_trans = np.array(transforms.Resize(rev_size)(to_pil(f_1[0])))
f_1_crf = crf_refine(np.array(batch["r_img"][0]), f_1_trans)
new_image = Image.new('RGB', (3 * image1_size[0], image1_size[1]),
(250, 250, 250))
img_res = Image.fromarray(f_1_crf)
new_image.paste(batch["r_img"][0], (0, 0))
new_image.paste(batch["r_mask"][0], (image1_size[0], 0))
new_image.paste(img_res, (image1_size[0] * 2, 0))
# The number of validation itteration
self.val_iter += 1
new_image.save(
os.path.join(self.args.msd_results_root, "Training",
"Eval_Epoch: " + str(self.val_iter) +
" Eval.png"))
# tensorboard_logs = {'avg_val_loss': avg_loss}
# # use key 'log'
# return {'val_loss': avg_loss, 'log': tensorboard_logs}
self.log('avg_val_loss', avg_loss)
def main():
net = BASE7().cuda(device_ids[0])
if len(args['snapshot']) > 0:
print('Load snapshot {} for testing'.format(args['snapshot']))
net.load_state_dict(
torch.load(
os.path.join(ckpt_path, exp_name, args['snapshot'] + '.pth')))
print('Load {} succeed!'.format(
os.path.join(ckpt_path, exp_name, args['snapshot'] + '.pth')))
net.eval()
with torch.no_grad():
for name, root in to_test.items():
img_list = [
img_name
for img_name in os.listdir(os.path.join(root, 'image'))
if img_name.endswith('.jpg')
]
start = time.time()
for idx, img_name in enumerate(img_list):
print('predicting for {}: {:>4d} / {}'.format(
name, idx + 1, len(img_list)))
check_mkdir(
os.path.join(ckpt_path, exp_name,
'%s_%s' % (exp_name, args['snapshot'])))
img = Image.open(os.path.join(root, 'image', img_name))
w, h = img.size
img_var = Variable(img_transform(img).unsqueeze(0)).cuda()
f, b, o = net(img_var)
foreground = f.data.squeeze(0).cpu()
background = b.data.squeeze(0).cpu()
output = o.data.squeeze(0).cpu()
prediction_f = np.array(
transforms.Resize((h, w))(to_pil(foreground)))
prediction_b = np.array(
transforms.Resize((h, w))(to_pil(background)))
prediction_o = np.array(
transforms.Resize((h, w))(to_pil(output)))
if args['crf']:
prediction_f = crf_refine(np.array(img.convert('RGB')),
prediction_f)
prediction_b = crf_refine(np.array(img.convert('RGB')),
prediction_b)
prediction_o = crf_refine(np.array(img.convert('RGB')),
prediction_o)
Image.fromarray(prediction_f).save(
os.path.join(ckpt_path, exp_name,
'%s_%s' % (exp_name, args['snapshot']),
img_name[:-4] + "_f.png"))
Image.fromarray(prediction_b).save(
os.path.join(ckpt_path, exp_name,
'%s_%s' % (exp_name, args['snapshot']),
img_name[:-4] + "_b.png"))
Image.fromarray(prediction_o).save(
os.path.join(ckpt_path, exp_name,
'%s_%s' % (exp_name, args['snapshot']),
img_name[:-4] + "_o.png"))
end = time.time()
print("Average Time Is : {:.2f}".format(
(end - start) / len(img_list)))
def main():
net = R3Net_prior(motion='GRU', se_layer=False, st_fuse=False)
print('load snapshot \'%s\' for testing' % args['snapshot'])
net.load_state_dict(
torch.load(os.path.join(ckpt_path, exp_name,
args['snapshot'] + '.pth'),
map_location='cuda:0'))
# net = train_online(net)
results = {}
for name, root in to_test.items():
precision_record, recall_record, = [AvgMeter() for _ in range(256)], [
AvgMeter() for _ in range(256)
]
mae_record = AvgMeter()
if args['save_results']:
check_mkdir(
os.path.join(ckpt_path, exp_name, '(%s) %s_%s' %
(exp_name, name, args['snapshot'])))
folders = os.listdir(root)
folders.sort()
for folder in folders:
net = train_online(net, seq_name=folder)
with torch.no_grad():
net.eval()
net.cuda()
imgs = os.listdir(os.path.join(root, folder))
imgs.sort()
for i in range(1, len(imgs) - args['batch_size'] + 1):
print(imgs[i])
img_var = []
img_names = []
for j in range(0, args['batch_size']):
img = Image.open(
os.path.join(root, folder,
imgs[i + j])).convert('RGB')
img_names.append(imgs[i + j])
shape = img.size
img = img.resize(args['input_size'])
img_var.append(
Variable(img_transform(img).unsqueeze(0),
volatile=True).cuda())
img_var = torch.cat(img_var, dim=0)
prediction = net(img_var)
precision = to_pil(prediction.data.squeeze(0).cpu())
precision = precision.resize(shape)
prediction = np.array(precision)
if args['crf_refine']:
prediction = crf_refine(np.array(img), prediction)
gt = np.array(
Image.open(
os.path.join(gt_root, folder, img_names[-1][:-4] +
'.png')).convert('L'))
precision, recall, mae = cal_precision_recall_mae(
prediction, gt)
for pidx, pdata in enumerate(zip(precision, recall)):
p, r = pdata
precision_record[pidx].update(p)
recall_record[pidx].update(r)
mae_record.update(mae)
if args['save_results']:
# folder, sub_name = os.path.split(img_names[-1])
save_path = os.path.join(
ckpt_path, exp_name,
'(%s) %s_%s' % (exp_name, name, args['snapshot']),
folder)
if not os.path.exists(save_path):
os.makedirs(save_path)
Image.fromarray(prediction).save(
os.path.join(save_path,
img_names[-1][:-4] + '.png'))
fmeasure = cal_fmeasure([precord.avg for precord in precision_record],
[rrecord.avg for rrecord in recall_record])
results[name] = {'fmeasure': fmeasure, 'mae': mae_record.avg}
print('test results:')
print(results)
def main():
##############################
# Using one GPU args.device_ids = [0]
##############################
if len(args.device_ids) == 1:
net = MirrorNet().cuda(device_ids[0])
optimizer = optim.Adam(net.parameters(), lr=args.lr, betas=args.betas)
if args.load_model:
# print(os.path.join(args.root_path + args.ckpt_path, args.exp_name, args.snapshot + '.pth'))
print('Load snapshot {} for testing'.format(args.snapshot))
net.load_state_dict(
torch.load(
os.path.join(args.ckpt_path, args.exp_name,
args.snapshot + '.pth')))
# net.load_state_dict(torch.load(os.path.join(args.ckpt_path, args.exp_name, args.snapshot + '.pth')))
print('Load {} succeed!'.format(
os.path.join(args.ckpt_path, args.exp_name,
args.snapshot + '.pth')))
if not args.train:
net.eval()
data_path = args.msd_testing_root
else:
data_path = args.msd_training_root
eval_path = args.msd_eval_root
net.train()
if args.developer_mode:
# To include the real images and masks
dataset = ImageFolder(data_path,
img_transform=img_transform,
target_transform=mask_transform,
add_real_imgs=True)
eval_dataset = ImageFolder(eval_path,
img_transform=img_transform,
target_transform=mask_transform,
add_real_imgs=True)
else:
dataset = ImageFolder(data_path,
img_transform=img_transform,
target_transform=mask_transform)
eval_dataset = ImageFolder(eval_path,
img_transform=img_transform,
target_transform=mask_transform)
loader = DataLoader(dataset,
batch_size=args.batch_size,
shuffle=args.shuffle_dataset)
eval_loader = DataLoader(eval_dataset, batch_size=1, shuffle=False)
# batch = dataset.sample(3)
# batch["img"][0].show()
# batch["mask"][0].show()
# print(batch)
if args.train:
print("Training")
idx = 0
##############################
# Training for number of epoches
for epoch in range(args.epochs):
start_time = time.time()
loss = 0
###############################
## Training
###############################
# Defining the tqdm progress bar for training dataset
pbar = tqdm(loader, desc="Processing batch number")
idx = 0
net.train()
for batch in pbar:
batch = dataset.sample(args.batch_size)
inputs = batch["img"]
outputs = batch["mask"]
inputs = torch.from_numpy(inputs)
outputs = torch.tensor(outputs)
# To GPU if available
if args.cuda:
inputs = inputs.cuda(device_ids[0])
outputs = outputs.cuda(device_ids[0])
# Getting the 4 different outputs
inputs.requires_grad = True
outputs.requires_grad = True
f_4_gpu, f_3_gpu, f_2_gpu, f_1_gpu = net(inputs)
if args.developer_mode:
if idx == 0:
f_1 = f_1_gpu.data.cpu()
rev_size = [
batch["size"][0][1], batch["size"][0][0]
]
image1_size = batch["size"][0]
f_1_trans = np.array(
transforms.Resize(rev_size)(to_pil(f_1[0])))
f_1_crf = crf_refine(np.array(batch["r_img"][0]),
f_1_trans)
new_image = Image.new(
'RGB', (3 * image1_size[0], image1_size[1]),
(250, 250, 250))
img_res = Image.fromarray(f_1_crf)
new_image.paste(batch["r_img"][0], (0, 0))
new_image.paste(batch["r_mask"][0],
(image1_size[0], 0))
new_image.paste(img_res, (image1_size[0] * 2, 0))
new_image.save(
os.path.join(
args.msd_results_root, "Testing",
"Epoch: " + str(epoch) + " Trining.png"))
print("Image saved")
idx += 1
##############################
# # For image processing
# f_4 = f_4_gpu.data.cpu()
# f_3 = f_3_gpu.data.cpu()
# f_2 = f_2_gpu.data.cpu()
# f_1 = f_1_gpu.data.cpu()
# f_4_arr = []
# f_3_arr = []
# f_2_arr = []
# f_1_arr = []
# f_1_arr_no_resize = []
# for i in range(args.batch_size):
# rev_size = [batch["size"][i][1], batch["size"][i][0]]
# f_4_arr.append(np.array(transforms.Resize(rev_size)(to_pil(f_4[i]))))
# f_3_arr.append(np.array(transforms.Resize(rev_size)(to_pil(f_3[i]))))
# f_2_arr.append(np.array(transforms.Resize(rev_size)(to_pil(f_2[i]))))
# f_1_arr.append(np.array(transforms.Resize(rev_size)(to_pil(f_1[i]))))
# f_1_arr_no_resize.append(np.array(to_pil(f_1[i])))
# if args.crf and args.developer_mode:
# f_1_arr[i] = crf_refine(np.array(batch["r_img"][i]), f_1_arr[i])
# # img_ = np.array(batch["img"][i])
# # img_ =img_.astype('uint8')
# f_1_arr_no_resize[i] = crf_refine(np.array(batch["r_img"][i].resize((args.scale, args.scale))), f_1_arr_no_resize[i])
# image1_size = batch["size"][i]
# if args.save_images:
# new_image = Image.new('RGB',(3*image1_size[0], image1_size[1]), (250,250,250))
# img_res = Image.fromarray(f_1_arr[i])
# new_image.paste(batch["r_img"][i],(0,0))
# new_image.paste(batch["r_mask"][i],(image1_size[0],0))
# new_image.paste(img_res,(image1_size[0]*2,0))
# new_image.save(os.path.join(args.msd_results_root, "Testing",
# "MNet_" + str(idx) +".png"))
# idx +=1
# img_res1 = Image.fromarray(f_2_arr[i])
# img_res2 = Image.fromarray(f_3_arr[i])
# img_res3 = Image.fromarray(f_4_arr[i])
# new_image.paste(img_res1,(0,0))
# new_image.paste(img_res2,(image1_size[0],0))
# new_image.paste(img_res3,(image1_size[0]*2,0))
# new_image.save(os.path.join(args.msd_results_root, "Testing",
# "MNet_" + str(idx) +".png"))
# idx +=1
# TODO
## Achieve loss for images with crf refining
## Starting Point
# f_1_arr = np.array(f_1_arr)
# f_1_arr_no_resize = np.array(f_1_arr_no_resize)
# f_1_arr_no_resize = torch.tensor(f_1_arr_no_resize)
# loss = lovasz_hinge(torch.tensor(batch["mask"]), f_1_arr_no_resize, per_image=False)
loss1 = lovasz_hinge(f_1_gpu, outputs,
per_image=False) * args.w_losses[0]
loss2 = lovasz_hinge(f_2_gpu, outputs,
per_image=False) * args.w_losses[1]
loss3 = lovasz_hinge(f_3_gpu, outputs,
per_image=False) * args.w_losses[2]
loss4 = lovasz_hinge(f_4_gpu, outputs,
per_image=False) * args.w_losses[3]
loss = loss1 + loss2 + loss3 + loss4
# outputs.requires_grad = False
# L2 = torch.nn.BCELoss()
# loss2 = L2(f_1_gpu, outputs)
# loss = loss2 +loss1
optimizer.zero_grad()
loss.backward()
optimizer.step()
res = loss.data.cpu()
text = "epoch: {}, Loss: {}".format(epoch, res)
text = text.ljust(40)
# text = text + " L_BCE: {}".format(loss2)
# text = text.ljust(60)
pbar.set_description(text)
print("Needed time:")
print(time.time() - start_time)
###############################
## Evaluation
###############################
# Defining the tqdm progress bar for evaluation dataset
eval_pbar = tqdm(eval_loader, desc="Processing batch number")
idx = 0
net.eval()
for batch in eval_pbar:
batch = dataset.sample(args.batch_size)
inputs = batch["img"]
outputs = batch["mask"]
inputs = torch.from_numpy(inputs)
outputs = torch.tensor(outputs)
# To GPU if available
if args.cuda:
inputs = inputs.cuda(device_ids[0])
outputs = outputs.cuda(device_ids[0])
# Getting the 4 different outputs
inputs.requires_grad = True
outputs.requires_grad = True
f_4_gpu, f_3_gpu, f_2_gpu, f_1_gpu = net(inputs)
if args.developer_mode:
if idx == 0:
f_1 = f_1_gpu.data.cpu()
rev_size = [
batch["size"][0][1], batch["size"][0][0]
]
image1_size = batch["size"][0]
f_1_trans = np.array(
transforms.Resize(rev_size)(to_pil(f_1[0])))
f_1_crf = crf_refine(np.array(batch["r_img"][0]),
f_1_trans)
new_image = Image.new(
'RGB', (3 * image1_size[0], image1_size[1]),
(250, 250, 250))
img_res = Image.fromarray(f_1_crf)
new_image.paste(batch["r_img"][0], (0, 0))
new_image.paste(batch["r_mask"][0],
(image1_size[0], 0))
new_image.paste(img_res, (image1_size[0] * 2, 0))
new_image.save(
os.path.join(
args.msd_results_root, "Testing",
"Epoch: " + str(epoch) + " Eval.png"))
print("Image saved")
idx += 1
eval_loss = lovasz_hinge(f_1_gpu, outputs, per_image=False)
res = eval_loss.data.cpu()
text = "epoch: {}, Eval Loss: {}".format(epoch, res)
text = text.ljust(45)
eval_pbar.set_description(text)
##############################
# Using multiple GPUs args.device_ids = [0, 1, ...]
##############################
else:
net = LitMirrorNet(args)
# net = net.load_from_checkpoint(args.root_path + args.ckpt_path + "/MirrorNet-epoch=16-val_loss=3.99.ckpt")
if args.load_model:
net = LitMirrorNet.load_from_checkpoint(args.ckpt_path +
args.ckpt_name,
args=args)
print('Loading {} checkpoint.'.format(args.ckpt_path +
args.ckpt_name))
trainer = Trainer(gpus=args.device_ids,
fast_dev_run=args.fast_dev_run,
accelerator='dp',
max_epochs=args.epochs,
callbacks=[checkpoint_callback],
check_val_every_n_epoch=args.val_every,
logger=tb_logger,
resume_from_checkpoint=args.ckpt_path +
args.ckpt_name)
print("Checkpoint loaded successfully!")
else:
trainer = Trainer(gpus=args.device_ids,
fast_dev_run=args.fast_dev_run,
accelerator='dp',
max_epochs=args.epochs,
callbacks=[checkpoint_callback],
check_val_every_n_epoch=args.val_every,
logger=tb_logger)
# resume_from_checkpoint = args.root_path + args.ckpt_path + "/MirrorNet-epoch=16-val_loss=3.99.ckpt")
if args.train:
print("Training")
trainer.fit(net)
final_epoch_model_path = args.ckpt_path + "final_epoch.ckpt"
trainer.save_checkpoint(final_epoch_model_path)
print("Done")
else:
print("Testing")
# trainer.test(model = net,
# ckpt_path = args.ckpt_path+args.ckpt_name)
trainer.test(model=net)
imgs.sort()
for img in imgs:
print(os.path.join(folder, img))
if name == 'VOS' or name == 'DAVSOD':
image = Image.open(os.path.join(root, folder,
img[:-4] + '.png')).convert('RGB')
else:
image = Image.open(os.path.join(root, folder,
img[:-4] + '.jpg')).convert('RGB')
gt = np.array(
Image.open(os.path.join(gt_root, folder, img)).convert('L'))
pred = np.array(
Image.open(os.path.join(save_path, folder, img)).convert('L'))
if args['crf_refine']:
pred = crf_refine(np.array(image), pred)
precision, recall, mae = cal_precision_recall_mae(pred, gt)
for pidx, pdata in enumerate(zip(precision, recall)):
p, r = pdata
precision_record[pidx].update(p)
recall_record[pidx].update(r)
mae_record.update(mae)
fmeasure = cal_fmeasure([precord.avg for precord in precision_record],
[rrecord.avg for rrecord in recall_record])
results[name] = {'fmeasure': fmeasure, 'mae': mae_record.avg}
print('test results:')
print(results)
def main():
net = AADFNet().cuda()
net = nn.DataParallel(net, device_ids=[0])
print exp_name + 'crf: ' + str(args['crf_refine'])
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():
results = {}
for name, root in to_test.iteritems():
precision_record, recall_record, = [
AvgMeter() for _ in range(256)
], [AvgMeter() for _ in range(256)]
mae_record = AvgMeter()
time_record = AvgMeter()
img_list = [
os.path.splitext(f)[0] for f in os.listdir(root)
if f.endswith('.jpg')
]
for idx, img_name in enumerate(img_list):
img_name = img_list[idx]
print 'predicting for %s: %d / %d' % (name, idx + 1,
len(img_list))
check_mkdir(
os.path.join(
ckpt_path, exp_name,
'(%s) %s_%s' % (exp_name, name, args['snapshot'])))
start = time.time()
img = Image.open(os.path.join(root, img_name +
'.jpg')).convert('RGB')
img_var = Variable(img_transform(img).unsqueeze(0),
volatile=True).cuda()
prediction = net(img_var)
W, H = img.size
prediction = F.upsample_bilinear(prediction, size=(H, W))
prediction = np.array(to_pil(prediction.data.squeeze(0).cpu()))
if args['crf_refine']:
prediction = crf_refine(np.array(img), prediction)
end = time.time()
gt = np.array(
Image.open(os.path.join(root,
img_name + '.png')).convert('L'))
precision, recall, mae = cal_precision_recall_mae(
prediction, gt)
for pidx, pdata in enumerate(zip(precision, recall)):
p, r = pdata
precision_record[pidx].update(p)
recall_record[pidx].update(r)
mae_record.update(mae)
time_record.update(end - start)
if args['save_results']:
Image.fromarray(prediction).save(
os.path.join(
ckpt_path, exp_name,
'(%s) %s_%s' % (exp_name, name, args['snapshot']),
img_name + '.png'))
max_fmeasure, mean_fmeasure = cal_fmeasure_both(
[precord.avg for precord in precision_record],
[rrecord.avg for rrecord in recall_record])
results[name] = {
'max_fmeasure': max_fmeasure,
'mae': mae_record.avg,
'mean_fmeasure': mean_fmeasure
}
print 'test results:'
print results
with open('Result', 'a') as f:
if args['crf_refine']:
f.write('with CRF')
f.write('Runing time %.6f \n' % time_record.avg)
f.write('\n%s\n iter:%s\n' % (exp_name, args['snapshot']))
for name, value in results.iteritems():
f.write(
'%s: mean_fmeasure: %.10f, mae: %.10f, max_fmeasure: %.10f\n'
% (name, value['mean_fmeasure'], value['mae'],
value['max_fmeasure']))
def main():
net = MirrorNet().cuda(device_ids[0])
if len(args['snapshot']) > 0:
print('Load snapshot {} for testing'.format(args['snapshot']))
# net.load_state_dict(torch.load(os.path.join(ckpt_path, exp_name, 'MirrorNet.pth')))
# print('Load {} succeed!'.format(os.path.join(ckpt_path, exp_name, 'MirrorNet.pth')))
net.load_state_dict(
torch.load(
os.path.join(ckpt_path, exp_name, args['snapshot'] + '.pth')))
print('Load {} succeed!'.format(
os.path.join(ckpt_path, exp_name, args['snapshot'] + '.pth')))
net.eval()
with torch.no_grad():
for name, root in to_test.items():
img_list = [
img_name
for img_name in os.listdir(os.path.join(root, 'image'))
]
# msk_list = [img_name for img_name in os.listdir(os.path.join(root, 'mask'))]
start = time.time()
for idx, img_name in enumerate(img_list):
msk_name = img_name.split('.')[0]
msk_name += ".png"
print('predicting for {}: {:>4d} / {}'.format(
name, idx + 1, len(img_list)))
check_mkdir(
os.path.join(
ckpt_path, exp_name,
'%s_%s_%s' % (exp_name, args['snapshot'], 'nocrf')))
img = Image.open(os.path.join(root, 'image', img_name))
msk = Image.open(os.path.join(root, 'mask', msk_name))
if img.mode != 'RGB':
img = img.convert('RGB')
print("{} is a gray image.".format(name))
w, h = img.size
img_var = Variable(img_transform(img).unsqueeze(0)).cuda(
device_ids[0])
f_4, f_3, f_2, f_1 = net(img_var)
f_4 = f_4.data.squeeze(0).cpu()
f_3 = f_3.data.squeeze(0).cpu()
f_2 = f_2.data.squeeze(0).cpu()
f_1 = f_1.data.squeeze(0).cpu()
f_4 = np.array(transforms.Resize((h, w))(to_pil(f_4)))
f_3 = np.array(transforms.Resize((h, w))(to_pil(f_3)))
f_2 = np.array(transforms.Resize((h, w))(to_pil(f_2)))
f_1 = np.array(transforms.Resize((h, w))(to_pil(f_1)))
if args['crf']:
f_1 = crf_refine(np.array(img.convert('RGB')), f_1)
Image.fromarray(f_1).save(
os.path.join(
ckpt_path, exp_name,
'%s_%s_%s' % (exp_name, args['snapshot'], 'nocrf'),
img_name[:-4] + ".png"))
#################################### My Addition ####################################
image1_size = img.size
new_image = Image.new('RGB',
(3 * image1_size[0], image1_size[1]),
(250, 250, 250))
img_res = Image.fromarray(f_1)
new_image.paste(img, (0, 0))
new_image.paste(msk, (image1_size[0], 0))
new_image.paste(img_res, (image1_size[0] * 2, 0))
# new_image.save(os.path.join(msd_results_root, '%s_%s' % (exp_name, args['snapshot']),
# img_name[:-4] + "_both" +".png"))
new_image.save(
os.path.join(msd_results_root, exp_name,
"MNet_" + str(idx) + ".png"))
#####################################################################################
end = time.time()
print("Average Time Is : {:.2f}".format(
(end - start) / len(img_list)))
def main():
net = R3Net(motion='')
print ('load snapshot \'%s\' for testing' % args['snapshot'])
net.load_state_dict(torch.load(os.path.join(ckpt_path, exp_name, args['snapshot'] + '.pth'), map_location='cuda:0'))
net.eval()
net.cuda()
results = {}
with torch.no_grad():
for name, root in to_test.items():
precision_record, recall_record, = [AvgMeter() for _ in range(256)], [AvgMeter() for _ in range(256)]
mae_record = AvgMeter()
if args['save_results']:
check_mkdir(os.path.join(ckpt_path, exp_name, '(%s) %s_%s' % (exp_name, name, args['snapshot'])))
img_list = [i_id.strip() for i_id in open(imgs_path)]
# img_list = [os.path.splitext(f)[0] for f in os.listdir(root) if f.endswith('.jpg')]
for idx, img_names in enumerate(img_list):
print ('predicting for %s: %d / %d' % (name, idx + 1, len(img_list)))
img_seq = img_names.split(',')
img_var = []
for img_name in img_seq:
img = Image.open(os.path.join(root, img_name + '.jpg')).convert('RGB')
shape = img.size
img = img.resize(args['input_size'])
img_var.append(Variable(img_transform(img).unsqueeze(0), volatile=True).cuda())
img_var = torch.cat(img_var, dim=0)
prediction = net(img_var)
precision = to_pil(prediction.data[-1].cpu())
precision = precision.resize(shape)
prediction = np.array(precision)
if args['crf_refine']:
prediction = crf_refine(np.array(img), prediction)
gt = np.array(Image.open(os.path.join(gt_root, img_seq[-1] + '.png')).convert('L'))
precision, recall, mae = cal_precision_recall_mae(prediction, gt)
for pidx, pdata in enumerate(zip(precision, recall)):
p, r = pdata
precision_record[pidx].update(p)
recall_record[pidx].update(r)
mae_record.update(mae)
if args['save_results']:
folder, sub_name = os.path.split(img_name)
save_path = os.path.join(ckpt_path, exp_name, '(%s) %s_%s' % (exp_name, name, args['snapshot']), folder)
if not os.path.exists(save_path):
os.makedirs(save_path)
Image.fromarray(prediction).save(os.path.join(save_path, sub_name + '.png'))
fmeasure = cal_fmeasure([precord.avg for precord in precision_record],
[rrecord.avg for rrecord in recall_record])
results[name] = {'fmeasure': fmeasure, 'mae': mae_record.avg}
print ('test results:')
print (results)
import numpy as np
import cv2
import skimage.io
from misc import crf_refine
DATA_DIR = "/home/iccd/data/msd9/train"
IMAGE_DIR = os.path.join(DATA_DIR, "image")
MASK_DIR = os.path.join(DATA_DIR, "mask")
EDGE_DIR = os.path.join(DATA_DIR, "edge")
if not os.path.exists(EDGE_DIR):
os.mkdir(EDGE_DIR)
masklist = os.listdir(MASK_DIR)
print("Total {} masks will be extracted edge!".format(len(masklist)))
for i, maskname in enumerate(masklist):
image_path = os.path.join(IMAGE_DIR, maskname[:-4] + '.jpg')
mask_path = os.path.join(MASK_DIR, maskname)
edge_path = os.path.join(EDGE_DIR, maskname)
image = skimage.io.imread(image_path)
mask = skimage.io.imread(mask_path)
edge = cv2.Canny(mask, 0, 255)
edge = crf_refine(image, edge)
# edge = np.where(edge != 0, 255, 0).astype(np.uint8)
cv2.imwrite(edge_path, edge)
print("{} {}".format(i, edge_path))
def main():
if args['model'] == 'BASNet':
net = BASNet(3, 1)
elif args['model'] == 'R3Net':
net = R3Net()
elif args['model'] == 'DSSNet':
net = build_model()
elif args['model'] == 'CPD':
net = CPD_ResNet()
elif args['model'] == 'RAS':
net = RAS()
elif args['model'] == 'PiCANet':
net = Unet()
elif args['model'] == 'PoolNet':
net = build_model_poolnet(base_model_cfg='resnet')
elif args['model'] == 'R2Net':
net = build_model_r2net(base_model_cfg='resnet')
elif args['model'] == 'F3Net':
net = F3Net(cfg=None)
print('load snapshot \'%s\' for testing' % args['snapshot'])
net.load_state_dict(
torch.load(os.path.join(ckpt_path, exp_name,
args['snapshot'] + '.pth'),
map_location='cuda:0'))
# net = train_online(net)
results = {}
for name, root in to_test.items():
precision_record, recall_record, = [AvgMeter() for _ in range(256)], [
AvgMeter() for _ in range(256)
]
mae_record = AvgMeter()
if args['save_results']:
check_mkdir(
os.path.join(
ckpt_path, exp_name, '(%s) %s_%s' %
(exp_name, name, args['snapshot'] + '_online')))
folders = os.listdir(root)
folders.sort()
folders = ['bmx-trees']
for folder in folders:
if args['online_train']:
net = train_online(net, seq_name=folder)
net.load_state_dict(
torch.load(os.path.join(
ckpt_path, exp_name,
str(args['snapshot']) + '_' + folder + '_online.pth'),
map_location='cuda:0'))
with torch.no_grad():
net.eval()
net.cuda()
imgs = os.listdir(os.path.join(root, folder))
imgs.sort()
for i in range(args['start'], len(imgs)):
print(imgs[i])
start = time.time()
img = Image.open(os.path.join(root, folder,
imgs[i])).convert('RGB')
shape = img.size
img = img.resize(args['input_size'])
img_var = Variable(img_transform(img).unsqueeze(0),
volatile=True).cuda()
if args['model'] == 'BASNet':
prediction, _, _, _, _, _, _, _ = net(img_var)
prediction = torch.sigmoid(prediction)
elif args['model'] == 'R3Net':
prediction = net(img_var)
elif args['model'] == 'DSSNet':
select = [1, 2, 3, 6]
prediction = net(img_var)
prediction = torch.mean(torch.cat(
[torch.sigmoid(prediction[i]) for i in select],
dim=1),
dim=1,
keepdim=True)
elif args['model'] == 'CPD':
_, prediction = net(img_var)
prediction = torch.sigmoid(prediction)
elif args['model'] == 'RAS':
prediction, _, _, _, _ = net(img_var)
prediction = torch.sigmoid(prediction)
elif args['model'] == 'PoolNet':
prediction = net(img_var)
prediction = torch.sigmoid(prediction)
elif args['model'] == 'F3Net':
_, prediction, _, _, _, _ = net(img_var)
prediction = torch.sigmoid(prediction)
elif args['model'] == 'R2Net':
_, _, _, _, _, prediction = net(img_var)
prediction = torch.sigmoid(prediction)
end = time.time()
print('running time:', (end - start))
if args['crf_refine']:
prediction = crf_refine(np.array(img), prediction)
precision = to_pil(prediction.data.squeeze(0).cpu())
precision = precision.resize(shape)
prediction = np.array(precision)
prediction = prediction.astype('float')
prediction = MaxMinNormalization(
prediction, prediction.max(), prediction.min()) * 255.0
prediction = prediction.astype('uint8')
gt = np.array(
Image.open(
os.path.join(gt_root, folder,
imgs[i][:-4] + '.png')).convert('L'))
precision, recall, mae = cal_precision_recall_mae(
prediction, gt)
for pidx, pdata in enumerate(zip(precision, recall)):
p, r = pdata
precision_record[pidx].update(p)
recall_record[pidx].update(r)
mae_record.update(mae)
if args['save_results']:
# folder, sub_name = os.path.split(imgs[i])
save_path = os.path.join(
ckpt_path, exp_name, '(%s) %s_%s' %
(exp_name, name, args['snapshot'] + '_online'),
folder)
if not os.path.exists(save_path):
os.makedirs(save_path)
Image.fromarray(prediction).save(
os.path.join(save_path, imgs[i]))
fmeasure = cal_fmeasure([precord.avg for precord in precision_record],
[rrecord.avg for rrecord in recall_record])
results[name] = {'fmeasure': fmeasure, 'mae': mae_record.avg}
print('test results:')
print(results)
def main():
net = R3Net_prior(motion='GRU',
se_layer=False,
attention=False,
pre_attention=True,
basic_model='resnet50',
sta=False,
naive_fuse=False)
print('load snapshot \'%s\' for testing' % args['snapshot'])
net.load_state_dict(
torch.load(os.path.join(ckpt_path, exp_name,
args['snapshot'] + '.pth'),
map_location='cuda:1'))
net.eval()
net.cuda()
results = {}
with torch.no_grad():
for name, root in to_test.items():
precision_record, recall_record, = [
AvgMeter() for _ in range(256)
], [AvgMeter() for _ in range(256)]
mae_record = AvgMeter()
if args['save_results']:
check_mkdir(
os.path.join(
ckpt_path, exp_name,
'(%s) %s_%s' % (exp_name, name, args['snapshot'])))
img_list = [i_id.strip() for i_id in open(imgs_path)]
# img_list = [os.path.splitext(f)[0] for f in os.listdir(root) if f.endswith('.jpg')]
for idx, img_names in enumerate(img_list):
print('predicting for %s: %d / %d' %
(name, idx + 1, len(img_list)))
img_seq = img_names.split(',')
for ratio in args['scale_ratio']:
prediction_scale = []
img_var = []
for img_name in img_seq:
if name == 'VOS' or name == 'DAVSOD':
img = Image.open(
os.path.join(root,
img_name + '.png')).convert('RGB')
else:
img = Image.open(
os.path.join(root,
img_name + '.jpg')).convert('RGB')
shape = img.size
new_dims = (int(img.size[0] * ratio),
int(img.size[1] * ratio))
img = img.resize(new_dims, Image.BILINEAR)
img = img.resize(args['input_size'])
img_var.append(
Variable(img_transform(img).unsqueeze(0),
volatile=True).cuda())
img_var = torch.cat(img_var, dim=0)
start = time.time()
prediction = net(img_var)
end = time.time()
print('running time:', (end - start))
prediction_scale.append(prediction)
prediction = torch.mean(torch.stack(prediction_scale, dim=0),
0)
prediction = to_pil(prediction.data.squeeze(0).cpu())
prediction = prediction.resize(shape)
prediction = np.array(prediction)
prediction = prediction.astype('float')
prediction = MaxMinNormalization(prediction, prediction.max(),
prediction.min()) * 255.0
prediction = prediction.astype('uint8')
if args['crf_refine']:
prediction = crf_refine(np.array(img), prediction)
gt = np.array(
Image.open(os.path.join(gt_root, img_seq[-1] +
'.png')).convert('L'))
precision, recall, mae = cal_precision_recall_mae(
prediction, gt)
for pidx, pdata in enumerate(zip(precision, recall)):
p, r = pdata
precision_record[pidx].update(p)
recall_record[pidx].update(r)
mae_record.update(mae)
if args['save_results']:
folder, sub_name = os.path.split(img_name)
save_path = os.path.join(
ckpt_path, exp_name,
'(%s) %s_%s' % (exp_name, name, args['snapshot']),
folder)
if not os.path.exists(save_path):
os.makedirs(save_path)
Image.fromarray(prediction).save(
os.path.join(save_path, sub_name + '.png'))
fmeasure = cal_fmeasure(
[precord.avg for precord in precision_record],
[rrecord.avg for rrecord in recall_record])
results[name] = {'fmeasure': fmeasure, 'mae': mae_record.avg}
print('test results:')
print(results)
def main():
# net = R3Net(motion='', se_layer=False, dilation=False, basic_model='resnet50')
net = Distill(basic_model='resnet50', seq=True)
print ('load snapshot \'%s\' for testing' % args['snapshot'])
net.load_state_dict(torch.load(os.path.join(ckpt_path, exp_name, args['snapshot'] + '.pth'), map_location='cuda:2'))
net.eval()
net.cuda()
results = {}
with torch.no_grad():
for name, root in to_test.items():
precision_record, recall_record, = [AvgMeter() for _ in range(256)], [AvgMeter() for _ in range(256)]
mae_record = AvgMeter()
if args['save_results']:
check_mkdir(os.path.join(ckpt_path, exp_name, '(%s) %s_%s' % (exp_name, name, args['snapshot'])))
img_list = [i_id.strip() for i_id in open(imgs_path)]
video = ''
for idx, img_name in enumerate(img_list):
print ('predicting for %s: %d / %d' % (name, idx + 1, len(img_list)))
print(img_name)
# img_var = []
if args['seq']:
if video != img_name.split('/')[0]:
video = img_name.split('/')[0]
if name == 'VOS' or name == 'DAVSOD':
img = Image.open(os.path.join(root, img_name + '.png')).convert('RGB')
else:
img = Image.open(os.path.join(root, img_name + '.jpg')).convert('RGB')
shape = img.size
img = img.resize(args['input_size'])
img_var = Variable(img_transform(img).unsqueeze(0), volatile=True).cuda()
start = time.time()
_, _, prediction = net(img_var, img_var, flag='seq')
end = time.time()
print('running time:', (end - start))
else:
if name == 'VOS' or name == 'DAVSOD':
img = Image.open(os.path.join(root, img_name + '.png')).convert('RGB')
pre = Image.open(os.path.join(root, img_list[idx - 1] + '.png')).convert('RGB')
else:
img = Image.open(os.path.join(root, img_name + '.jpg')).convert('RGB')
pre = Image.open(os.path.join(root, img_list[idx - 1] + '.jpg')).convert('RGB')
shape = img.size
img = img.resize(args['input_size'])
pre = pre.resize(args['input_size'])
img_var = Variable(img_transform(img).unsqueeze(0), volatile=True).cuda()
pre_var = Variable(img_transform(pre).unsqueeze(0), volatile=True).cuda()
start = time.time()
_, _, prediction = net(pre_var, img_var, flag='seq')
end = time.time()
print('running time:', (end - start))
else:
start = time.time()
if name == 'VOS' or name == 'DAVSOD':
img = Image.open(os.path.join(root, img_name + '.png')).convert('RGB')
else:
img = Image.open(os.path.join(root, img_name + '.jpg')).convert('RGB')
shape = img.size
img = img.resize(args['input_size'])
img_var = Variable(img_transform(img).unsqueeze(0), volatile=True).cuda()
_, _, prediction, _ = net(img_var, img_var)
end = time.time()
print('running time:', (end - start))
precision = to_pil(prediction.data.squeeze(0).cpu())
precision = precision.resize(shape)
prediction = np.array(precision)
prediction = prediction.astype('float')
prediction = MaxMinNormalization(prediction, prediction.max(), prediction.min()) * 255.0
prediction = prediction.astype('uint8')
if args['crf_refine']:
prediction = crf_refine(np.array(img), prediction)
gt = np.array(Image.open(os.path.join(gt_root, img_name + '.png')).convert('L'))
precision, recall, mae = cal_precision_recall_mae(prediction, gt)
for pidx, pdata in enumerate(zip(precision, recall)):
p, r = pdata
precision_record[pidx].update(p)
recall_record[pidx].update(r)
mae_record.update(mae)
if args['save_results']:
folder, sub_name = os.path.split(img_name)
save_path = os.path.join(ckpt_path, exp_name, '(%s) %s_%s' % (exp_name, name, args['snapshot']), folder)
if not os.path.exists(save_path):
os.makedirs(save_path)
Image.fromarray(prediction).save(os.path.join(save_path, sub_name + '.png'))
fmeasure = cal_fmeasure([precord.avg for precord in precision_record],
[rrecord.avg for rrecord in recall_record])
results[name] = {'fmeasure': fmeasure, 'mae': mae_record.avg}
print ('test results:')
print (results)