def forward(self, data_input, model_output):
if self.flownet is None:
self._setup()
targets = data_input['targets'].to(device)
outputs = model_output['outputs'].to(device)
flows = self.flownet.infer_video(targets).to(device)
from utils.flow_utils import warp_optical_flow
warped_targets = warp_optical_flow(targets[:, :-1], -flows).detach()
warped_outputs = warp_optical_flow(outputs[:, :-1], -flows).detach()
non_occlusion_masks = self._get_non_occlusion_mask(
targets, warped_targets)
# model_output is passed by name and dictionary is mutable
# These values are sent to trainer for visualization
model_output['warped_outputs'] = warped_outputs[0]
model_output['warped_targets'] = warped_targets[0]
model_output['non_occlusion_masks'] = non_occlusion_masks[0]
from utils.flow_utils import flow_to_image
flow_imgs = []
for flow in flows[0]:
flow_img = flow_to_image(flow.cpu().permute(
1, 2, 0).detach().numpy()).transpose(2, 0, 1)
flow_imgs.append(torch.Tensor(flow_img))
model_output['flow_imgs'] = flow_imgs
masks = data_input['masks'].to(device)
return self._get_loss(outputs, warped_outputs, non_occlusion_masks,
masks)
'--img_list',
nargs='+',
default=['examples/img1.png', 'examples/img2.png'])
args = parser.parse_args()
cfg = {
'model': {
'upsample': True,
'n_frames': len(args.img_list),
'reduce_dense': True
},
'pretrained_model': args.model,
'test_shape': args.test_shape,
}
ts = TestHelper(cfg)
imgs = [imageio.imread(img).astype(np.float32) for img in args.img_list]
h, w = imgs[0].shape[:2]
flow_12 = ts.run(imgs)['flows_fw'][0]
flow_12 = resize_flow(flow_12, (h, w))
np_flow_12 = flow_12[0].detach().cpu().numpy().transpose([1, 2, 0])
vis_flow = flow_to_image(np_flow_12)
fig = plt.figure()
plt.imshow(vis_flow)
plt.show()
# vx.astype(np.float32).tofile(f)
# f.close()
#
# f = open(os.path.join(r'G:\ARFlow-master\data\flow_dataset\ceshi_tmp', name + '_vy.bin'), 'wb')
# vy.astype(np.float32).tofile(f)
# f.close()
occu_mask_12 = occu_mask1[0].detach().cpu().numpy().transpose(
[1, 2, 0])
occu_mask_21 = occu_mask2[0].detach().cpu().numpy().transpose(
[1, 2, 0])
back_occu_mask_1 = back_occu_mask1[0].detach().cpu().numpy().transpose(
[1, 2, 0])
back_occu_mask_2 = back_occu_mask2[0].detach().cpu().numpy().transpose(
[1, 2, 0])
vis_flow_12 = flow_to_image(np_flow_12)
vis_flow_21 = flow_to_image(np_flow_21)
row1 = np.concatenate([
np.uint8(imgs[0]),
np.uint8(imgs[1]),
np.uint8(vis_flow_12),
np.uint8(vis_flow_21)
],
axis=1)
# row2 = np.concatenate([np.uint8(np_warped_image12), np.uint8(np_warped_image21), 255*np.tile(occu_mask_12, (1, 1, 3)),
# 255*np.tile(occu_mask_21, (1, 1, 3))], axis=1)
row2 = np.concatenate([
np.uint8(np_warped_image12),
np.uint8(np_warped_image21), 255 * np.tile(back_occu_mask_1,
(1, 1, 3)),