def main():
teacher = MGA_Network(nInputChannels=3,
n_classes=1,
os=16,
img_backbone_type='resnet101',
flow_backbone_type='resnet34')
teacher = load_MGA(teacher, args['mga_model_path'], device_id=device_id)
teacher.eval()
teacher.cuda(device_id2)
student = Ensemble(device_id).cuda(device_id).train()
# fix_parameters(net.named_parameters())
optimizer = optim.SGD([{
'params': [
param for name, param in student.named_parameters()
if name[-4:] == 'bias'
],
'lr':
2 * args['lr']
}, {
'params': [
param for name, param in student.named_parameters()
if name[-4:] != 'bias'
],
'lr':
args['lr'],
'weight_decay':
args['weight_decay']
}],
momentum=args['momentum'])
if len(args['snapshot']) > 0:
print('training resumes from ' + args['snapshot'])
student.load_state_dict(
torch.load(
os.path.join(ckpt_path, exp_name, args['snapshot'] + '.pth')))
optimizer.load_state_dict(
torch.load(
os.path.join(ckpt_path, exp_name,
args['snapshot'] + '_optim.pth')))
optimizer.param_groups[0]['lr'] = 2 * args['lr']
optimizer.param_groups[1]['lr'] = args['lr']
if len(args['pretrain']) > 0:
print('pretrain model from ' + args['pretrain'])
student = load_part_of_model(student,
args['pretrain'],
device_id=device_id)
# fix_parameters(student.named_parameters())
check_mkdir(ckpt_path)
check_mkdir(os.path.join(ckpt_path, exp_name))
open(log_path, 'w').write(str(args) + '\n\n')
train(student, teacher, optimizer)
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')