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 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)
results = {}
# save_path = os.path.join(ckpt_path, exp_name, '(%s) %s_%s' % (exp_name, name, args['snapshot']))
folders = os.listdir(root_inference)
folders.sort()
for folder in folders:
imgs = os.listdir(os.path.join(root_inference, folder))
imgs.sort()
for img in imgs:
print(os.path.join(folder, img))
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(root_inference, folder, img)).convert('L'))
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 = 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)
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 = Ensemble(device_id, pretrained=False)
print ('load snapshot \'%s\' for testing' % args['snapshot'])
# net.load_state_dict(torch.load('pretrained/R2Net.pth', map_location='cuda:2'))
# net = load_part_of_model2(net, 'pretrained/R2Net.pth', device_id=2)
net.load_state_dict(torch.load(os.path.join(ckpt_path, exp_name, args['snapshot'] + '.pth'),
map_location='cuda:' + str(device_id)))
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)]
for idx, img_name in enumerate(img_list):
print('predicting for %s: %d / %d' % (name, idx + 1, len(img_list)))
print(img_name)
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()
outputs_a, outputs_c = net(img_var)
a_out1u, a_out2u, a_out2r, a_out3r, a_out4r, a_out5r = outputs_a # F3Net
# b_outputs0, b_outputs1 = outputs_b # CPD
c_outputs0, c_outputs1, c_outputs2, c_outputs3, c_outputs4 = outputs_c # RAS
prediction = torch.sigmoid(c_outputs0)
end = time.time()
print('running time:', (end - start))
# e = Erosion2d(1, 1, 5, soft_max=False).cuda()
# prediction2 = e(prediction)
#
# precision2 = to_pil(prediction2.data.squeeze(0).cpu())
# precision2 = prediction2.data.squeeze(0).cpu().numpy()
# precision2 = precision2.resize(shape)
# prediction2 = np.array(precision2)
# prediction2 = prediction2.astype('float')
precision = to_pil(prediction.data.squeeze(0).cpu())
precision = precision.resize(shape)
prediction = np.array(precision)
prediction = prediction.astype('float')
# plt.style.use('classic')
# plt.subplot(1, 2, 1)
# plt.imshow(prediction)
# plt.subplot(1, 2, 2)
# plt.imshow(precision2[0])
# plt.show()
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)
log_path = os.path.join('result_all.txt')
open(log_path, 'a').write(exp_name + ' ' + args['snapshot'] + '\n')
open(log_path, 'a').write(str(results) + '\n\n')
def main():
net = ConcatNet(backbone='resnet50', embedding=128, batch_mode='old')
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():
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, 'DAVIS', args['snapshot'])))
index = 0
old_seq = ''
# img_list = [os.path.splitext(f)[0] for f in os.listdir(root) if f.endswith('.jpg')]
for idx, batch in enumerate(testloader):
print('%d processd ' % (idx))
target = batch['target']
seq_name = batch['seq_name']
args['seq_name'] = seq_name[0]
print('sequence name:', seq_name[0])
if old_seq == args['seq_name']:
index = index + 1
else:
index = 0
output_sum = 0
for i in range(0, args['sample_range']):
search = batch['search_' + str(i)]
output = net(
Variable(target, volatile=True).cuda(),
Variable(search, volatile=True).cuda())
output = F.upsample(output,
size=target.size()[2:],
mode='bilinear',
align_corners=True)
output = F.sigmoid(output)
output_sum = output_sum + output[0].data.cpu().numpy()
output_final = output_sum / args['sample_range']
output_final = cv2.resize(output_final[0], (854, 480))
output_final = (output_final * 255).astype(np.uint8)
gt = np.array(
Image.open(
os.path.join(args['gt_dir'], args['seq_name'],
str(index).zfill(5) + '.png')).convert('L'))
precision, recall, mae = cal_precision_recall_mae(output_final, 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']:
old_seq = args['seq_name']
save_path = os.path.join(
ckpt_path, exp_name,
'(%s) %s_%s' % (exp_name, 'DAVIS', args['snapshot']),
args['seq_name'])
if not os.path.exists(save_path):
os.makedirs(save_path)
Image.fromarray(output_final).save(
os.path.join(save_path,
str(index).zfill(5) + '.png'))
fmeasure = cal_fmeasure([precord.avg for precord in precision_record],
[rrecord.avg for rrecord in recall_record])
results['DAVIS'] = {'fmeasure': fmeasure, 'mae': mae_record.avg}
print('test results:')
print(results)
def main():
net = R3Net(num_class=1).cuda().train()
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:0'}))
# when using the linux compute then run the code
state_dict = torch.load(
os.path.join(ckpt_path, exp_name, args['snapshot'] + '.pth'))
# new_state_dict = OrderedDict()
# for k, v in state_dict.items():
# namekey = k[7:] # remove 'module.'
# new_state_dict[namekey] = v
# # load params
# net.load_state_dict(new_state_dict)
net.load_state_dict(state_dict)
net.eval()
results = {}
with torch.no_grad():
for name, root in to_test.items():
pre, rec = [], []
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'])))
hha = sorted(glob.glob(root + '/hha/*.jpg'))
# hha = sorted(glob.glob(root + '/depth/*.png'))
# # load nju
# rgb = ['./data/NJU_test/LR/'+i.split('/')[4] for i in sorted(glob.glob(root + '/hha/*.jpg'))]
# label = ['./data/NJU_test/GT/'+i.split('/')[4].split('.')[0]+'.png' for i in sorted(glob.glob(root + '/hha/*.jpg'))]
# # load nlpr
rgb = [
'./data/NLPR_test/RGB/' + i.split('/')[4].split('_Depth')[0] +
'.jpg' for i in sorted(glob.glob(root + '/hha/*.jpg'))
]
label = [
'./data/NLPR_test/groundtruth/' +
i.split('/')[4].split('_Depth')[0] + '.jpg'
for i in sorted(glob.glob(root + '/hha/*.jpg'))
]
# # load des
# rgb = ['./data/DES/image/RGBD_data_'+i.split('_')[2]+'.png' for i in sorted(glob.glob(root + '/hha/*.jpg'))]
# label = ['./data/DES/GT/RGBD_data_'+i.split('_')[2]+'.bmp' for i in sorted(glob.glob(root + '/hha/*.jpg'))]
# load lfsd
# rgb = ['./data/LFSD/all_focus_images/'+i.split('/')[4] for i in sorted(glob.glob(root + '/hha/*.jpg'))]
# label =['./data/LFSD/ground_truth/'+i.split('/')[4].split('.')[0]+'.png' for i in sorted(glob.glob(root + '/hha/*.jpg'))]
# load ssd
# rgb = ['./data/SSD100/RGB/' + i.split('/')[4] for i in sorted(glob.glob(root + '/hha/*.jpg'))]
# label = ['./data/SSD100/GT/' + i.split('/')[4].split('.')[0] + '.png' for i in sorted(glob.glob(root + '/hha/*.jpg'))]
# load sip
# rgb = ['./data/SIP/RGB/' + i.split('/')[4] for i in sorted(glob.glob(root + '/hha/*.jpg'))]
# label = ['./data/SIP/GT/' + i.split('/')[4].split('.')[0] + '.png' for i in sorted(glob.glob(root + '/hha/*.jpg'))]
# load stere
# rgb = ['./data/STERE/RGB/' + i.split('/')[4] for i in sorted(glob.glob(root + '/hha/*.jpg'))]
# label = ['./data/STERE/GT/' + i.split('/')[4].split('.')[0] + '.png' for i in sorted(glob.glob(root + '/hha/*.jpg'))]
# hha = ['./data/STERE/depth/' + i.split('/')[4].split('.')[0] + '.png' for i in sorted(glob.glob(root + '/hha/*.jpg'))]
for idx in range(len(rgb)):
print('predicting for %s: %d / %d' % (name, idx + 1, len(rgb)))
# deal lfsd,sip and stere dataset
img = Image.open(rgb[idx]).convert('RGB')
depth = Image.open(hha[idx]).convert('RGB')
# depth = cv2.imread(hha[idx])
# depth = torch.cat((depth,depth,depth),2)
# depth = Image.fromarray(cv2.cvtColor(depth,cv2.COLOR_BGR2RGB))
h, w = img.size[0], img.size[1]
img = img_transform(img)
depth = img_transform(depth)
img_var = Variable(img.unsqueeze(0), volatile=True).cuda()
depth_var = Variable(depth.unsqueeze(0), volatile=True).cuda()
# prediction = net(img_var,depth_var)#,out3,out2,out1,out0
prediction = net(img_var, depth_var)
#for i in range(len(tmp)-1):
# tmp_ = np.array(to_pil(tmp[i+1].data.squeeze(0).cpu()))
# Image.fromarray(tmp_).save('depth_visual/'+str(i)+'.jpg')
prediction = F.interpolate(input=prediction,
size=(w, h),
mode='bilinear',
align_corners=False)
prediction = np.array(to_pil(prediction.data.squeeze(0).cpu()))
# load GT
gt = np.array(Image.open(label[idx]).convert('L'))
try:
precision, recall, mae = cal_precision_recall_mae(
prediction, gt)
pre.append(precision)
rec.append(recall)
except Exception as error:
import pdb
pdb.set_trace()
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']), rgb[idx].split('/')[4].split('.')[0]+'.png'))
# if args['save_results']:
# cv2.imwrite(os.path.join(ckpt_path, exp_name, '(%s) %s_%s' % (exp_name, name, args['snapshot']), rgb[idx].split('/')[4]),prediction)
max_fmeasure, mean_fmeasure = cal_fmeasure1(pre, rec)
# max_fmeasure,mean_fmeasure = cal_fmeasure([precord.avg for precord in precision_record], [rrecord.avg for rrecord in recall_record])
results[name] = {
'max_fmeasure': max_fmeasure,
'mean_fmeasure': mean_fmeasure,
'mae': mae_record.avg
}
print('test results:')
print(results)
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)
def main():
# net = R3Net(motion='', se_layer=False, dilation=False, basic_model='resnet50')
net = SNet(cfg=None)
print('load snapshot \'%s\' for testing' % args['snapshot'])
# net.load_state_dict(torch.load('pretrained/R2Net.pth', map_location='cuda:2'))
# net = load_part_of_model2(net, 'pretrained/R2Net.pth', device_id=2)
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 = ''
pre_predict = None
for idx, img_name in enumerate(img_list):
print('predicting for %s: %d / %d' %
(name, idx + 1, len(img_list)))
print(img_name)
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()
if args['model'] == 'BASNet':
prediction, _, prediction2, _, _, _, _, _ = 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':
prediction2, 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':
prediction2, prediction, _, _, _, _ = net(img_var)
prediction = torch.sigmoid(prediction)
elif args['model'] == 'R2Net':
_, _, _, _, _, prediction = net(img_var)
prediction = torch.sigmoid(prediction)
end = time.time()
pre_predict = prediction
print('running time:', (end - start))
else:
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)
end = time.time()
print('running time:', (end - start))
pre_predict = prediction
# e = Erosion2d(1, 1, 5, soft_max=False).cuda()
# prediction2 = e(prediction)
#
# precision2 = to_pil(prediction2.data.squeeze(0).cpu())
# precision2 = prediction2.data.squeeze(0).cpu().numpy()
# precision2 = precision2.resize(shape)
# prediction2 = np.array(precision2)
# prediction2 = prediction2.astype('float')
precision = to_pil(prediction.data.squeeze(0).cpu())
precision = precision.resize(shape)
prediction = np.array(precision)
prediction = prediction.astype('float')
# plt.style.use('classic')
# plt.subplot(1, 2, 1)
# plt.imshow(prediction)
# plt.subplot(1, 2, 2)
# plt.imshow(precision2[0])
# plt.show()
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)
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():
net = SDCNet(num_classes=5).cuda()
print('load snapshot \'%s\' for testing, mode:\'%s\'' %
(args['snapshot'], args['test_mode']))
print(exp_name)
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.items():
print('load snapshot \'%s\' for testing %s' %
(args['snapshot'], name))
test_data = pd.read_csv(root)
test_set = TestFolder_joint(test_data, joint_transform,
img_transform, target_transform)
test_loader = DataLoader(test_set,
batch_size=1,
num_workers=0,
shuffle=False)
precision0_record, recall0_record, = [
AvgMeter() for _ in range(256)
], [AvgMeter() for _ in range(256)]
precision1_record, recall1_record, = [
AvgMeter() for _ in range(256)
], [AvgMeter() for _ in range(256)]
precision2_record, recall2_record, = [
AvgMeter() for _ in range(256)
], [AvgMeter() for _ in range(256)]
precision3_record, recall3_record, = [
AvgMeter() for _ in range(256)
], [AvgMeter() for _ in range(256)]
precision4_record, recall4_record, = [
AvgMeter() for _ in range(256)
], [AvgMeter() for _ in range(256)]
precision5_record, recall5_record, = [
AvgMeter() for _ in range(256)
], [AvgMeter() for _ in range(256)]
precision6_record, recall6_record, = [
AvgMeter() for _ in range(256)
], [AvgMeter() for _ in range(256)]
mae0_record = AvgMeter()
mae1_record = AvgMeter()
mae2_record = AvgMeter()
mae3_record = AvgMeter()
mae4_record = AvgMeter()
mae5_record = AvgMeter()
mae6_record = AvgMeter()
n0, n1, n2, n3, n4, n5 = 0, 0, 0, 0, 0, 0
if args['save_results']:
check_mkdir(
os.path.join(ckpt_path, exp_name,
'%s_%s' % (name, args['snapshot'])))
for i, (inputs, gt, labels,
img_path) in enumerate(tqdm(test_loader)):
shape = gt.size()[2:]
img_var = Variable(inputs).cuda()
img = np.array(to_pil(img_var.data.squeeze(0).cpu()))
gt = np.array(to_pil(gt.data.squeeze(0).cpu()))
sizec = labels.numpy()
pred2021 = net(img_var, sizec)
pred2021 = F.interpolate(pred2021,
size=shape,
mode='bilinear',
align_corners=True)
pred2021 = np.array(to_pil(pred2021.data.squeeze(0).cpu()))
if labels == 0:
precision1, recall1, mae1 = cal_precision_recall_mae(
pred2021, gt)
for pidx, pdata in enumerate(zip(precision1, recall1)):
p, r = pdata
precision1_record[pidx].update(p)
#print('Presicion:', p, 'Recall:', r)
recall1_record[pidx].update(r)
mae1_record.update(mae1)
n1 += 1
elif labels == 1:
precision2, recall2, mae2 = cal_precision_recall_mae(
pred2021, gt)
for pidx, pdata in enumerate(zip(precision2, recall2)):
p, r = pdata
precision2_record[pidx].update(p)
#print('Presicion:', p, 'Recall:', r)
recall2_record[pidx].update(r)
mae2_record.update(mae2)
n2 += 1
elif labels == 2:
precision3, recall3, mae3 = cal_precision_recall_mae(
pred2021, gt)
for pidx, pdata in enumerate(zip(precision3, recall3)):
p, r = pdata
precision3_record[pidx].update(p)
#print('Presicion:', p, 'Recall:', r)
recall3_record[pidx].update(r)
mae3_record.update(mae3)
n3 += 1
elif labels == 3:
precision4, recall4, mae4 = cal_precision_recall_mae(
pred2021, gt)
for pidx, pdata in enumerate(zip(precision4, recall4)):
p, r = pdata
precision4_record[pidx].update(p)
#print('Presicion:', p, 'Recall:', r)
recall4_record[pidx].update(r)
mae4_record.update(mae4)
n4 += 1
elif labels == 4:
precision5, recall5, mae5 = cal_precision_recall_mae(
pred2021, gt)
for pidx, pdata in enumerate(zip(precision5, recall5)):
p, r = pdata
precision5_record[pidx].update(p)
#print('Presicion:', p, 'Recall:', r)
recall5_record[pidx].update(r)
mae5_record.update(mae5)
n5 += 1
precision6, recall6, mae6 = cal_precision_recall_mae(
pred2021, gt)
for pidx, pdata in enumerate(zip(precision6, recall6)):
p, r = pdata
precision6_record[pidx].update(p)
recall6_record[pidx].update(r)
mae6_record.update(mae6)
img_name = os.path.split(str(img_path))[1]
img_name = os.path.splitext(img_name)[0]
n0 += 1
if args['save_results']:
Image.fromarray(pred2021).save(
os.path.join(ckpt_path, exp_name,
'%s_%s' % (name, args['snapshot']),
img_name + '_2021.png'))
fmeasure1 = cal_fmeasure(
[precord.avg for precord in precision1_record],
[rrecord.avg for rrecord in recall1_record])
fmeasure2 = cal_fmeasure(
[precord.avg for precord in precision2_record],
[rrecord.avg for rrecord in recall2_record])
fmeasure3 = cal_fmeasure(
[precord.avg for precord in precision3_record],
[rrecord.avg for rrecord in recall3_record])
fmeasure4 = cal_fmeasure(
[precord.avg for precord in precision4_record],
[rrecord.avg for rrecord in recall4_record])
fmeasure5 = cal_fmeasure(
[precord.avg for precord in precision5_record],
[rrecord.avg for rrecord in recall5_record])
fmeasure6 = cal_fmeasure(
[precord.avg for precord in precision6_record],
[rrecord.avg for rrecord in recall6_record])
results[name] = {
'fmeasure1': fmeasure1,
'mae1': mae1_record.avg,
'fmeasure2': fmeasure2,
'mae2': mae2_record.avg,
'fmeasure3': fmeasure3,
'mae3': mae3_record.avg,
'fmeasure4': fmeasure4,
'mae4': mae4_record.avg,
'fmeasure5': fmeasure5,
'mae5': mae5_record.avg,
'fmeasure6': fmeasure6,
'mae6': mae6_record.avg
}
print('test results:')
print('[fmeasure1 %.3f], [mae1 %.4f], [class1 %.0f]\n'\
'[fmeasure2 %.3f], [mae2 %.4f], [class2 %.0f]\n'\
'[fmeasure3 %.3f], [mae3 %.4f], [class3 %.0f]\n'\
'[fmeasure4 %.3f], [mae4 %.4f], [class4 %.0f]\n'\
'[fmeasure5 %.3f], [mae5 %.4f], [class5 %.0f]\n'\
'[fmeasure6 %.3f], [mae6 %.4f], [all %.0f]\n'%\
(fmeasure1, mae1_record.avg, n1, fmeasure2, mae2_record.avg, n2, fmeasure3, mae3_record.avg, n3, fmeasure4, mae4_record.avg, n4, fmeasure5, mae5_record.avg, n5, fmeasure6, mae6_record.avg, n0))
