def main():
net = AFENet(classes=19, pretrained_model_path=None).cuda()
net.load_state_dict(
torch.load(os.path.join(args['model_save_path'], args['snapshot'])))
net.eval()
mean_std = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
input_transform = standard_transforms.Compose([
standard_transforms.ToTensor(),
standard_transforms.Normalize(*mean_std)
])
img_path = args['image_path']
result_save_path = args['test_result_save_path']
scales = args['scale']
img = Image.open(img_path).convert('RGB')
crop_size = 768
prediction = np.zeros((19, 1024, 2048), dtype=float)
origin_w, origin_h = img.size
for scale in scales:
new_w = int(origin_w * scale)
new_h = int(origin_h * scale)
if scale == 1.0:
scale_img = img
else:
scale_img = img.resize((new_w, new_h), Image.BILINEAR)
prediction += scale_process(net, scale_img, origin_w, origin_h,
crop_size, new_w, new_h, input_transform)
prediction /= len(scales)
prediction = np.argmax(prediction, axis=0)
prediction = prediction.reshape([1024, 2048])
prediction = cityscapes.colorize_mask(prediction)
prediction.save(os.path.join(result_save_path))
print("Completed.")
def validate(t_val_loader, net, input_transform, restore, scales, count):
net.eval()
inputs_all, gts_all, predictions_all = [], [], []
crop_size = 768
for vi, (inputs, targets) in enumerate(t_val_loader):
prediction = np.zeros((19, 1024, 2048), dtype=float)
img = inputs.data[0]
img = restore(img)
inputs_all.append(img)
origin_w, origin_h = img.size
for scale in scales:
new_w = int(origin_w * scale)
new_h = int(origin_h * scale)
if scale == 1.0:
scale_img = img
else:
scale_img = img.resize((new_w, new_h), Image.BILINEAR)
prediction += scale_process(net, scale_img, origin_w, origin_h,
crop_size, new_w, new_h,
input_transform)
prediction /= len(scales)
prediction = np.argmax(prediction, axis=0)
prediction = prediction.reshape([1, 1024, 2048])
gts_all.append(targets.data.cpu().numpy())
predictions_all.append(prediction)
if (vi + 1) % 20 == 0:
localtime = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
print('%s Completed %d / %d' % (localtime, vi + 1, count))
gts_all = np.concatenate(gts_all)
predictions_all = np.concatenate(predictions_all)
print(gts_all.shape)
print(predictions_all.shape)
acc, acc_cls, mean_iu, fwavacc = evaluate(predictions_all, gts_all,
cityscapes.num_classes)
if args['val_save_or_not']:
check_mkdir(args['val_exp_save_path'])
for idx, data in enumerate(zip(inputs_all, gts_all, predictions_all)):
if data[0] is None:
continue
input_pil = data[0]
gt_pil = cityscapes.colorize_mask(data[1])
predictions_pil = cityscapes.colorize_mask(data[2])
input_pil.save(
os.path.join(args['val_exp_save_path'], '%d_input.png' % idx))
predictions_pil.save(
os.path.join(args['val_exp_save_path'],
'%d_prediction.png' % idx))
gt_pil.save(
os.path.join(args['val_exp_save_path'], '%d_gt.png' % idx))
print(
'-----------------------------------------------------------------------------------------------------------'
)
print('[acc %.5f], [acc_cls %.5f], [mean_iu %.5f], [fwavacc %.5f]' %
(acc, acc_cls, mean_iu, fwavacc))
print(
'-----------------------------------------------------------------------------------------------------------'
)
