def validate_flow_with_gt(val_loader, flow_net, epoch, logger, output_writers=[]): global args batch_time = AverageMeter() error_names = ['epe_total', 'epe_rigid', 'epe_non_rigid', 'outliers'] errors = AverageMeter(i=len(error_names)) log_outputs = len(output_writers) > 0 # switch to evaluate mode flow_net.eval() end = time.time() for i, (tgt_img, ref_imgs, intrinsics, intrinsics_inv, flow_gt, obj_map_gt) in enumerate(val_loader): tgt_img_var = Variable(tgt_img.cuda(), volatile=True) ref_imgs_var = [Variable(img.cuda(), volatile=True) for img in ref_imgs] flow_gt_var = Variable(flow_gt.cuda(), volatile=True) obj_map_gt_var = Variable(obj_map_gt.cuda(), volatile=True) if args.flownet == 'Back2Future': flow_fwd, flow_bwd= flow_net(tgt_img_var, ref_imgs_var) else: flow_fwd = flow_net(tgt_img_var, ref_imgs_var[1]) flow_bwd = flow_net(tgt_img_var, ref_imgs_var[0]) if args.DEBUG: flow_fwd_x = flow_fwd[:,0].view(-1).abs().data flow_gt_var_x = flow_gt_var[:,0].view(-1).abs().data flow_fwd_non_rigid = flow_fwd total_flow = flow_fwd obj_map_gt_var_expanded = obj_map_gt_var.unsqueeze(1).type_as(flow_fwd) if log_outputs and i % 10 == 0 and i/10 < len(output_writers): index = int(i//10) if epoch == 0: output_writers[index].add_image('val flow Input', tensor2array(tgt_img[0]), 0) flow_to_show = flow_gt[0][:2,:,:].cpu() output_writers[index].add_image('val target Flow', flow_to_image(tensor2array(flow_to_show)), epoch) output_writers[index].add_image('val Non-rigid Flow Output', flow_to_image(tensor2array(flow_fwd_non_rigid.data[0].cpu())), epoch) if np.isnan(flow_gt.sum().item()) or np.isnan(total_flow.data.sum().item()): print('NaN encountered') _epe_errors = compute_all_epes(flow_gt_var, flow_fwd, flow_fwd, (1-obj_map_gt_var_expanded) ) errors.update(_epe_errors) batch_time.update(time.time() - end) end = time.time() if args.log_terminal: logger.valid_bar.update(i) if i % args.print_freq == 0: logger.valid_writer.write('valid: Time {} Abs Error {:.4f} ({:.4f})'.format(batch_time, errors.val[0], errors.avg[0])) if args.log_terminal: logger.valid_bar.update(len(val_loader)) return errors.avg, error_names
def validate_flow_with_gt(patch, mask, patch_shape, val_loader, flow_net, epoch, logger, output_writer): global args batch_time = AverageMeter() error_names = ['epe', 'adv_epe', 'cos_sim', 'adv_cos_sim'] errors = AverageMeter(i=len(error_names)) flow_net.eval() end = time.time() for i, (ref_img_past, tgt_img, ref_img_future, flow_gt, _, _, _) in enumerate(val_loader): tgt_img_var = Variable(tgt_img.cuda(), volatile=True) ref_img_past_var = Variable(ref_img_past.cuda(), volatile=True) ref_img_future_var = Variable(ref_img_future.cuda(), volatile=True) flow_gt_var = Variable(flow_gt.cuda(), volatile=True) if type(flow_net).__name__ == 'Back2Future': flow_fwd = flow_net(ref_img_past_var, tgt_img_var, ref_img_future_var) else: flow_fwd = flow_net(tgt_img_var, ref_img_future_var) data_shape = tgt_img.cpu().numpy().shape if args.patch_type == 'circle': patch_full, mask_full, _, _, _, _ = circle_transform( patch, mask, patch.copy(), data_shape, patch_shape) elif args.patch_type == 'square': patch_full, mask_full, _, _, _ = square_transform( patch, mask, patch.copy(), data_shape, patch_shape, norotate=args.norotate) patch_full, mask_full = torch.FloatTensor( patch_full), torch.FloatTensor(mask_full) patch_full, mask_full = patch_full.cuda(), mask_full.cuda() patch_var, mask_var = Variable(patch_full), Variable(mask_full) adv_tgt_img_var = torch.mul( (1 - mask_var), tgt_img_var) + torch.mul(mask_var, patch_var) if type(flow_net).__name__ == 'Back2Future': adv_ref_img_past_var = torch.mul( (1 - mask_var), ref_img_past_var) + torch.mul( mask_var, patch_var) adv_ref_img_future_var = torch.mul( (1 - mask_var), ref_img_future_var) + torch.mul( mask_var, patch_var) adv_tgt_img_var = torch.clamp(adv_tgt_img_var, -1, 1) if type(flow_net).__name__ == 'Back2Future': adv_ref_img_past_var = torch.clamp(adv_ref_img_past_var, -1, 1) adv_ref_img_future_var = torch.clamp(adv_ref_img_future_var, -1, 1) if type(flow_net).__name__ == 'Back2Future': adv_flow_fwd = flow_net(adv_ref_img_past_var, adv_tgt_img_var, adv_ref_img_future_var) else: adv_flow_fwd = flow_net(adv_tgt_img_var, adv_ref_img_future_var) epe = compute_epe(gt=flow_gt_var, pred=flow_fwd) adv_epe = compute_epe(gt=flow_gt_var, pred=adv_flow_fwd) cos_sim = compute_cossim(flow_gt_var, flow_fwd) adv_cos_sim = compute_cossim(flow_gt_var, adv_flow_fwd) errors.update([epe, adv_epe, cos_sim, adv_cos_sim]) if args.log_output and i % 10 == 0: index = int(i // 10) if epoch == 0: output_writer.add_image( 'val flow Input', transpose_image(tensor2array(tgt_img[0])), 0) flow_to_show = flow_gt[0][:2, :, :].cpu() output_writer.add_image( 'val target Flow', transpose_image(flow_to_image(tensor2array(flow_to_show))), epoch) val_Flow_Output = transpose_image( flow_to_image(tensor2array(flow_fwd.data[0].cpu()))) / 255. val_adv_Flow_Output = transpose_image( flow_to_image(tensor2array(adv_flow_fwd.data[0].cpu()))) / 255. val_Diff_Flow_Output = transpose_image( flow_to_image( tensor2array( (adv_flow_fwd - flow_fwd).data[0].cpu()))) / 255. val_adv_tgt_image = transpose_image( tensor2array(adv_tgt_img_var.data.cpu()[0])) if type(flow_net).__name__ == 'Back2Future': val_adv_ref_past_image = transpose_image( tensor2array(adv_ref_img_past_var.data.cpu()[0])) val_adv_ref_future_image = transpose_image( tensor2array(adv_ref_img_future_var.data.cpu()[0])) val_patch = transpose_image(tensor2array(patch_var.data.cpu()[0])) if type(flow_net).__name__ == 'Back2Future': val_output_viz = np.hstack( (val_Flow_Output, val_adv_Flow_Output, val_Diff_Flow_Output, val_adv_ref_past_image, val_adv_tgt_image, val_adv_ref_future_image)) else: val_output_viz = np.hstack( (val_Flow_Output, val_adv_Flow_Output, val_Diff_Flow_Output, val_adv_tgt_image, val_adv_ref_future_image)) output_writer.add_image('val Output viz {}'.format(index), val_output_viz, epoch) if args.log_terminal: logger.valid_bar.update(i) batch_time.update(time.time() - end) end = time.time() return errors.avg, error_names
def main(): global args args = parser.parse_args() args.pretrained_disp = Path(args.pretrained_disp) args.pretrained_pose = Path(args.pretrained_pose) args.pretrained_mask = Path(args.pretrained_mask) args.pretrained_flow = Path(args.pretrained_flow) if args.output_dir is not None: args.output_dir = Path(args.output_dir) args.output_dir.makedirs_p() image_dir = args.output_dir / 'images' gt_dir = args.output_dir / 'gt' mask_dir = args.output_dir / 'mask' viz_dir = args.output_dir / 'viz' rigidity_mask_dir = args.output_dir / 'rigidity' rigidity_census_mask_dir = args.output_dir / 'rigidity_census' explainability_mask_dir = args.output_dir / 'explainability' image_dir.makedirs_p() gt_dir.makedirs_p() mask_dir.makedirs_p() viz_dir.makedirs_p() rigidity_mask_dir.makedirs_p() rigidity_census_mask_dir.makedirs_p() explainability_mask_dir.makedirs_p() output_writer = SummaryWriter(args.output_dir) normalize = custom_transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) flow_loader_h, flow_loader_w = 256, 832 valid_flow_transform = custom_transforms.Compose([ custom_transforms.Scale(h=flow_loader_h, w=flow_loader_w), custom_transforms.ArrayToTensor(), normalize ]) val_flow_set = ValidationMask(root=args.kitti_dir, sequence_length=5, transform=valid_flow_transform) val_loader = torch.utils.data.DataLoader(val_flow_set, batch_size=1, shuffle=False, num_workers=2, pin_memory=True, drop_last=True) disp_net = getattr(models, args.dispnet)().cuda() pose_net = getattr(models, args.posenet)(nb_ref_imgs=4).cuda() mask_net = getattr(models, args.masknet)(nb_ref_imgs=4).cuda() flow_net = getattr(models, args.flownet)(nlevels=args.nlevels).cuda() dispnet_weights = torch.load(args.pretrained_disp) posenet_weights = torch.load(args.pretrained_pose) masknet_weights = torch.load(args.pretrained_mask) flownet_weights = torch.load(args.pretrained_flow) disp_net.load_state_dict(dispnet_weights['state_dict']) pose_net.load_state_dict(posenet_weights['state_dict']) flow_net.load_state_dict(flownet_weights['state_dict']) mask_net.load_state_dict(masknet_weights['state_dict']) disp_net.eval() pose_net.eval() mask_net.eval() flow_net.eval() error_names = ['tp_0', 'fp_0', 'fn_0', 'tp_1', 'fp_1', 'fn_1'] errors = AverageMeter(i=len(error_names)) errors_census = AverageMeter(i=len(error_names)) errors_bare = AverageMeter(i=len(error_names)) for i, (tgt_img, ref_imgs, intrinsics, intrinsics_inv, flow_gt, obj_map_gt, semantic_map_gt) in enumerate(tqdm(val_loader)): tgt_img_var = Variable(tgt_img.cuda(), volatile=True) ref_imgs_var = [ Variable(img.cuda(), volatile=True) for img in ref_imgs ] intrinsics_var = Variable(intrinsics.cuda(), volatile=True) intrinsics_inv_var = Variable(intrinsics_inv.cuda(), volatile=True) flow_gt_var = Variable(flow_gt.cuda(), volatile=True) obj_map_gt_var = Variable(obj_map_gt.cuda(), volatile=True) disp = disp_net(tgt_img_var) depth = 1 / disp pose = pose_net(tgt_img_var, ref_imgs_var) explainability_mask = mask_net(tgt_img_var, ref_imgs_var) if args.flownet in ['Back2Future']: flow_fwd, flow_bwd, _ = flow_net(tgt_img_var, ref_imgs_var[1:3]) else: flow_fwd = flow_net(tgt_img_var, ref_imgs_var[2]) flow_cam = pose2flow(depth.squeeze(1), pose[:, 2], intrinsics_var, intrinsics_inv_var) rigidity_mask = 1 - (1 - explainability_mask[:, 1]) * ( 1 - explainability_mask[:, 2]).unsqueeze(1) > 0.5 rigidity_mask_census_soft = (flow_cam - flow_fwd).pow(2).sum( dim=1).unsqueeze(1).sqrt() #.normalize() rigidity_mask_census_soft = 1 - rigidity_mask_census_soft / rigidity_mask_census_soft.max( ) rigidity_mask_census = rigidity_mask_census_soft > args.THRESH rigidity_mask_combined = 1 - ( 1 - rigidity_mask.type_as(explainability_mask)) * ( 1 - rigidity_mask_census.type_as(explainability_mask)) flow_fwd_non_rigid = (1 - rigidity_mask_combined).type_as( flow_fwd).expand_as(flow_fwd) * flow_fwd flow_fwd_rigid = rigidity_mask_combined.type_as(flow_fwd).expand_as( flow_fwd) * flow_cam total_flow = flow_fwd_rigid + flow_fwd_non_rigid obj_map_gt_var_expanded = obj_map_gt_var.unsqueeze(1).type_as(flow_fwd) tgt_img_np = tgt_img[0].numpy() rigidity_mask_combined_np = rigidity_mask_combined.cpu().data[0].numpy( ) rigidity_mask_census_np = rigidity_mask_census.cpu().data[0].numpy() rigidity_mask_bare_np = rigidity_mask.cpu().data[0].numpy() gt_mask_np = obj_map_gt[0].numpy() semantic_map_np = semantic_map_gt[0].numpy() _errors = mask_error(gt_mask_np, semantic_map_np, rigidity_mask_combined_np[0]) _errors_census = mask_error(gt_mask_np, semantic_map_np, rigidity_mask_census_np[0]) _errors_bare = mask_error(gt_mask_np, semantic_map_np, rigidity_mask_bare_np[0]) errors.update(_errors) errors_census.update(_errors_census) errors_bare.update(_errors_bare) if args.output_dir is not None: np.save(image_dir / str(i).zfill(3), tgt_img_np) np.save(gt_dir / str(i).zfill(3), gt_mask_np) np.save(mask_dir / str(i).zfill(3), rigidity_mask_combined_np) np.save(rigidity_mask_dir / str(i).zfill(3), rigidity_mask.cpu().data[0].numpy()) np.save(rigidity_census_mask_dir / str(i).zfill(3), rigidity_mask_census.cpu().data[0].numpy()) np.save(explainability_mask_dir / str(i).zfill(3), explainability_mask[:, 1].cpu().data[0].numpy()) # rigidity_mask_dir rigidity_mask.numpy() # rigidity_census_mask_dir rigidity_mask_census.numpy() if (args.output_dir is not None) and i % 10 == 0: ind = int(i // 10) output_writer.add_image( 'val Dispnet Output Normalized', tensor2array(disp.data[0].cpu(), max_value=None, colormap='bone'), ind) output_writer.add_image('val Input', tensor2array(tgt_img[0].cpu()), i) output_writer.add_image( 'val Total Flow Output', flow_to_image(tensor2array(total_flow.data[0].cpu())), ind) output_writer.add_image( 'val Rigid Flow Output', flow_to_image(tensor2array(flow_fwd_rigid.data[0].cpu())), ind) output_writer.add_image( 'val Non-rigid Flow Output', flow_to_image(tensor2array(flow_fwd_non_rigid.data[0].cpu())), ind) output_writer.add_image( 'val Rigidity Mask', tensor2array(rigidity_mask.data[0].cpu(), max_value=1, colormap='bone'), ind) output_writer.add_image( 'val Rigidity Mask Census', tensor2array(rigidity_mask_census.data[0].cpu(), max_value=1, colormap='bone'), ind) output_writer.add_image( 'val Rigidity Mask Combined', tensor2array(rigidity_mask_combined.data[0].cpu(), max_value=1, colormap='bone'), ind) if args.output_dir is not None: tgt_img_viz = tensor2array(tgt_img[0].cpu()) depth_viz = tensor2array(disp.data[0].cpu(), max_value=None, colormap='magma') mask_viz = tensor2array(rigidity_mask_census_soft.data[0].cpu(), max_value=1, colormap='bone') row2_viz = flow_to_image( np.hstack((tensor2array(flow_cam.data[0].cpu()), tensor2array(flow_fwd_non_rigid.data[0].cpu()), tensor2array(total_flow.data[0].cpu())))) row1_viz = np.hstack((tgt_img_viz, depth_viz, mask_viz)) ####### sửa 2 cái vstack thành hstack ############### viz3 = np.hstack( (255 * tgt_img_viz, 255 * depth_viz, 255 * mask_viz, flow_to_image( np.hstack((tensor2array(flow_fwd_non_rigid.data[0].cpu()), tensor2array(total_flow.data[0].cpu())))))) ######################################################## ######## code tự thêm #################### row1_viz = np.transpose(row1_viz, (1, 2, 0)) row2_viz = np.transpose(row2_viz, (1, 2, 0)) viz3 = np.transpose(viz3, (1, 2, 0)) ########################################## row1_viz_im = Image.fromarray((255 * row1_viz).astype('uint8')) row2_viz_im = Image.fromarray((row2_viz).astype('uint8')) viz3_im = Image.fromarray(viz3.astype('uint8')) row1_viz_im.save(viz_dir / str(i).zfill(3) + '01.png') row2_viz_im.save(viz_dir / str(i).zfill(3) + '02.png') viz3_im.save(viz_dir / str(i).zfill(3) + '03.png') bg_iou = errors.sum[0] / (errors.sum[0] + errors.sum[1] + errors.sum[2]) fg_iou = errors.sum[3] / (errors.sum[3] + errors.sum[4] + errors.sum[5]) avg_iou = (bg_iou + fg_iou) / 2 bg_iou_census = errors_census.sum[0] / ( errors_census.sum[0] + errors_census.sum[1] + errors_census.sum[2]) fg_iou_census = errors_census.sum[3] / ( errors_census.sum[3] + errors_census.sum[4] + errors_census.sum[5]) avg_iou_census = (bg_iou_census + fg_iou_census) / 2 bg_iou_bare = errors_bare.sum[0] / ( errors_bare.sum[0] + errors_bare.sum[1] + errors_bare.sum[2]) fg_iou_bare = errors_bare.sum[3] / ( errors_bare.sum[3] + errors_bare.sum[4] + errors_bare.sum[5]) avg_iou_bare = (bg_iou_bare + fg_iou_bare) / 2 print("Results Full Model") print("\t {:>10}, {:>10}, {:>10} ".format('iou', 'bg_iou', 'fg_iou')) print("Errors \t {:10.4f}, {:10.4f} {:10.4f}".format( avg_iou, bg_iou, fg_iou)) print("Results Census only") print("\t {:>10}, {:>10}, {:>10} ".format('iou', 'bg_iou', 'fg_iou')) print("Errors \t {:10.4f}, {:10.4f} {:10.4f}".format( avg_iou_census, bg_iou_census, fg_iou_census)) print("Results Bare") print("\t {:>10}, {:>10}, {:>10} ".format('iou', 'bg_iou', 'fg_iou')) print("Errors \t {:10.4f}, {:10.4f} {:10.4f}".format( avg_iou_bare, bg_iou_bare, fg_iou_bare))
def main(): global args args = parser.parse_args() args.pretrained_disp = Path(args.pretrained_disp) args.pretrained_pose = Path(args.pretrained_pose) args.pretrained_mask = Path(args.pretrained_mask) args.pretrained_flow = Path(args.pretrained_flow) if args.output_dir is not None: args.output_dir = Path(args.output_dir) args.output_dir.makedirs_p() image_dir = args.output_dir / 'images' gt_dir = args.output_dir / 'gt' mask_dir = args.output_dir / 'mask' viz_dir = args.output_dir / 'viz' image_dir.makedirs_p() gt_dir.makedirs_p() mask_dir.makedirs_p() viz_dir.makedirs_p() output_writer = SummaryWriter(args.output_dir) normalize = custom_transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) flow_loader_h, flow_loader_w = 256, 832 valid_flow_transform = custom_transforms.Compose([ custom_transforms.Scale(h=flow_loader_h, w=flow_loader_w), custom_transforms.ArrayToTensor(), normalize ]) if args.dataset == "kitti2015": val_flow_set = ValidationFlow(root=args.kitti_dir, sequence_length=5, transform=valid_flow_transform) val_loader = torch.utils.data.DataLoader(val_flow_set, batch_size=1, shuffle=False, num_workers=2, pin_memory=True, drop_last=True) disp_net = getattr(models, args.dispnet)().cuda() pose_net = getattr(models, args.posenet)(nb_ref_imgs=4).cuda() mask_net = getattr(models, args.masknet)(nb_ref_imgs=4).cuda() flow_net = getattr(models, args.flownet)(nlevels=args.nlevels).cuda() dispnet_weights = torch.load(args.pretrained_disp) posenet_weights = torch.load(args.pretrained_pose) masknet_weights = torch.load(args.pretrained_mask) flownet_weights = torch.load(args.pretrained_flow) disp_net.load_state_dict(dispnet_weights['state_dict']) pose_net.load_state_dict(posenet_weights['state_dict']) flow_net.load_state_dict(flownet_weights['state_dict']) mask_net.load_state_dict(masknet_weights['state_dict']) disp_net.eval() pose_net.eval() mask_net.eval() flow_net.eval() error_names = [ 'epe_total', 'epe_sp', 'epe_mv', 'Fl', 'epe_total_gt_mask', 'epe_sp_gt_mask', 'epe_mv_gt_mask', 'Fl_gt_mask' ] errors = AverageMeter(i=len(error_names)) for i, (tgt_img, ref_imgs, intrinsics, intrinsics_inv, flow_gt, obj_map_gt) in enumerate(tqdm(val_loader)): tgt_img_var = Variable(tgt_img.cuda(), volatile=True) ref_imgs_var = [ Variable(img.cuda(), volatile=True) for img in ref_imgs ] intrinsics_var = Variable(intrinsics.cuda(), volatile=True) intrinsics_inv_var = Variable(intrinsics_inv.cuda(), volatile=True) flow_gt_var = Variable(flow_gt.cuda(), volatile=True) obj_map_gt_var = Variable(obj_map_gt.cuda(), volatile=True) disp = disp_net(tgt_img_var) depth = 1 / disp pose = pose_net(tgt_img_var, ref_imgs_var) explainability_mask = mask_net(tgt_img_var, ref_imgs_var) if args.flownet == 'Back2Future': flow_fwd, flow_bwd, _ = flow_net(tgt_img_var, ref_imgs_var[1:3]) else: flow_fwd = flow_net(tgt_img_var, ref_imgs_var[2]) flow_cam = pose2flow(depth.squeeze(1), pose[:, 2], intrinsics_var, intrinsics_inv_var) flow_cam_bwd = pose2flow(depth.squeeze(1), pose[:, 1], intrinsics_var, intrinsics_inv_var) rigidity_mask = 1 - (1 - explainability_mask[:, 1]) * ( 1 - explainability_mask[:, 2]).unsqueeze(1) > 0.5 rigidity_mask_census_soft = (flow_cam - flow_fwd).abs() #.normalize() rigidity_mask_census_u = rigidity_mask_census_soft[:, 0] < args.THRESH rigidity_mask_census_v = rigidity_mask_census_soft[:, 1] < args.THRESH rigidity_mask_census = (rigidity_mask_census_u).type_as(flow_fwd) * ( rigidity_mask_census_v).type_as(flow_fwd) rigidity_mask_combined = 1 - ( 1 - rigidity_mask.type_as(explainability_mask)) * ( 1 - rigidity_mask_census.type_as(explainability_mask)) obj_map_gt_var_expanded = obj_map_gt_var.unsqueeze(1).type_as(flow_fwd) flow_fwd_non_rigid = (rigidity_mask_combined <= args.THRESH).type_as( flow_fwd).expand_as(flow_fwd) * flow_fwd flow_fwd_rigid = (rigidity_mask_combined > args.THRESH ).type_as(flow_cam).expand_as(flow_cam) * flow_cam total_flow = flow_fwd_rigid + flow_fwd_non_rigid rigidity_mask = rigidity_mask.type_as(flow_fwd) _epe_errors = compute_all_epes( flow_gt_var, flow_cam, flow_fwd, rigidity_mask_combined) + compute_all_epes( flow_gt_var, flow_cam, flow_fwd, (1 - obj_map_gt_var_expanded)) errors.update(_epe_errors) tgt_img_np = tgt_img[0].numpy() rigidity_mask_combined_np = rigidity_mask_combined.cpu().data[0].numpy( ) gt_mask_np = obj_map_gt[0].numpy() if args.output_dir is not None: np.save(image_dir / str(i).zfill(3), tgt_img_np) np.save(gt_dir / str(i).zfill(3), gt_mask_np) np.save(mask_dir / str(i).zfill(3), rigidity_mask_combined_np) if (args.output_dir is not None) and i % 10 == 0: ind = int(i // 10) output_writer.add_image( 'val Dispnet Output Normalized', tensor2array(disp.data[0].cpu(), max_value=None, colormap='bone'), ind) output_writer.add_image('val Input', tensor2array(tgt_img[0].cpu()), i) output_writer.add_image( 'val Total Flow Output', flow_to_image(tensor2array(total_flow.data[0].cpu())), ind) output_writer.add_image( 'val Rigid Flow Output', flow_to_image(tensor2array(flow_fwd_rigid.data[0].cpu())), ind) output_writer.add_image( 'val Non-rigid Flow Output', flow_to_image(tensor2array(flow_fwd_non_rigid.data[0].cpu())), ind) output_writer.add_image( 'val Rigidity Mask', tensor2array(rigidity_mask.data[0].cpu(), max_value=1, colormap='bone'), ind) output_writer.add_image( 'val Rigidity Mask Census', tensor2array(rigidity_mask_census.data[0].cpu(), max_value=1, colormap='bone'), ind) output_writer.add_image( 'val Rigidity Mask Combined', tensor2array(rigidity_mask_combined.data[0].cpu(), max_value=1, colormap='bone'), ind) tgt_img_viz = tensor2array(tgt_img[0].cpu()) depth_viz = tensor2array(disp.data[0].cpu(), max_value=None, colormap='bone') mask_viz = tensor2array( rigidity_mask_census_soft.data[0].prod(dim=0).cpu(), max_value=1, colormap='bone') rigid_flow_viz = flow_to_image(tensor2array( flow_cam.data[0].cpu())) non_rigid_flow_viz = flow_to_image( tensor2array(flow_fwd_non_rigid.data[0].cpu())) total_flow_viz = flow_to_image( tensor2array(total_flow.data[0].cpu())) row1_viz = np.hstack((tgt_img_viz, depth_viz, mask_viz)) row2_viz = np.hstack( (rigid_flow_viz, non_rigid_flow_viz, total_flow_viz)) row1_viz_im = Image.fromarray((255 * row1_viz).astype('uint8')) row2_viz_im = Image.fromarray((row2_viz).astype('uint8')) row1_viz_im.save(viz_dir / str(i).zfill(3) + '01.png') row2_viz_im.save(viz_dir / str(i).zfill(3) + '02.png') print("Results") print("\t {:>10}, {:>10}, {:>10}, {:>6}, {:>10}, {:>10}, {:>10}, {:>10} ". format(*error_names)) print( "Errors \t {:10.4f}, {:10.4f}, {:10.4f}, {:10.4f}, {:10.4f}, {:10.4f}, {:10.4f}, {:10.4f}" .format(*errors.avg))
def train(train_loader, mask_net, pose_net, flow_net, optimizer, epoch_size, train_writer): global args, n_iter w1 = args.smooth_loss_weight w2 = args.mask_loss_weight w3 = args.consensus_loss_weight w4 = args.flow_loss_weight mask_net.train() pose_net.train() flow_net.train() average_loss = 0 for i, (rgb_tgt_img, rgb_ref_imgs, depth_tgt_img, depth_ref_imgs, intrinsics, intrinsics_inv, pose_list) in enumerate(tqdm(train_loader)): rgb_tgt_img_var = Variable(rgb_tgt_img.cuda()) # print(rgb_tgt_img_var.size()) rgb_ref_imgs_var = [Variable(img.cuda()) for img in rgb_ref_imgs] # rgb_ref_imgs_var = [rgb_ref_imgs_var[0], rgb_ref_imgs_var[0], rgb_ref_imgs_var[1], rgb_ref_imgs_var[1]] depth_tgt_img_var = Variable(depth_tgt_img.unsqueeze(1).cuda()) depth_ref_imgs_var = [ Variable(img.unsqueeze(1).cuda()) for img in depth_ref_imgs ] intrinsics_var = Variable(intrinsics.cuda()) intrinsics_inv_var = Variable(intrinsics_inv.cuda()) # pose_list_var = [Variable(one_pose.float().cuda()) for one_pose in pose_list] explainability_mask = mask_net(rgb_tgt_img_var, rgb_ref_imgs_var) valid_pixle_mask = torch.where( depth_tgt_img_var == 0, torch.zeros_like(depth_tgt_img_var), torch.ones_like(depth_tgt_img_var)) # zero is invalid # print(depth_test[0].sum()) # print(explainability_mask[0].size()) #torch.Size([4, 2, 384, 512]) # print() pose = pose_net(rgb_tgt_img_var, rgb_ref_imgs_var) # generate flow from camera pose and depth flow_fwd, flow_bwd, _ = flow_net(rgb_tgt_img_var, rgb_ref_imgs_var) flows_cam_fwd = pose2flow(depth_ref_imgs_var[1].squeeze(1), pose[:, 1], intrinsics_var, intrinsics_inv_var) flows_cam_bwd = pose2flow(depth_ref_imgs_var[0].squeeze(1), pose[:, 0], intrinsics_var, intrinsics_inv_var) rigidity_mask_fwd = (flows_cam_fwd - flow_fwd[0]).abs() rigidity_mask_bwd = (flows_cam_bwd - flow_bwd[0]).abs() # loss 1: smoothness loss loss1 = smooth_loss(explainability_mask) + smooth_loss( flow_bwd) + smooth_loss(flow_fwd) # loss 2: explainability loss loss2 = explainability_loss(explainability_mask) # loss 3 consensus loss (the mask from networks and the mask from residual) depth_Res_mask, depth_ref_img_warped, depth_diff = depth_residual_mask( valid_pixle_mask, explainability_mask[0], rgb_tgt_img_var, rgb_ref_imgs_var, intrinsics_var, intrinsics_inv_var, depth_tgt_img_var, pose) # print(depth_Res_mask[0].size(), explainability_mask[0].size()) loss3 = consensus_loss(explainability_mask[0], rigidity_mask_bwd, rigidity_mask_fwd, args.THRESH, args.wbce) # loss 4: flow loss loss4, flow_ref_img_warped, flow_diff = flow_loss( rgb_tgt_img_var, rgb_ref_imgs_var, [flow_bwd, flow_fwd], explainability_mask) # compute gradient and do Adam step loss = w1 * loss1 + w2 * loss2 + w3 * loss3 + w4 * loss4 average_loss += loss.item() optimizer.zero_grad() loss.backward() optimizer.step() # visualization in tensorboard if i > 0 and n_iter % args.print_freq == 0: train_writer.add_scalar('smoothness loss', loss1.item(), n_iter) train_writer.add_scalar('explainability loss', loss2.item(), n_iter) train_writer.add_scalar('consensus loss', loss3.item(), n_iter) train_writer.add_scalar('flow loss', loss4.item(), n_iter) train_writer.add_scalar('total loss', loss.item(), n_iter) if n_iter % (args.training_output_freq) == 0: train_writer.add_image('train Input', tensor2array(rgb_tgt_img_var[0]), n_iter) train_writer.add_image( 'train Exp mask Outputs ', tensor2array(explainability_mask[0][0, 0].data.cpu(), max_value=1, colormap='bone'), n_iter) train_writer.add_image( 'train depth Res mask ', tensor2array(depth_Res_mask[0][0].data.cpu(), max_value=1, colormap='bone'), n_iter) train_writer.add_image( 'train depth ', tensor2array(depth_tgt_img_var[0].data.cpu(), max_value=1, colormap='bone'), n_iter) train_writer.add_image( 'train valid pixel ', tensor2array(valid_pixle_mask[0].data.cpu(), max_value=1, colormap='bone'), n_iter) train_writer.add_image( 'train after mask', tensor2array(rgb_tgt_img_var[0] * explainability_mask[0][0, 0]), n_iter) train_writer.add_image('train depth diff', tensor2array(depth_diff[0]), n_iter) train_writer.add_image('train flow diff', tensor2array(flow_diff[0]), n_iter) train_writer.add_image('train depth warped img', tensor2array(depth_ref_img_warped[0]), n_iter) train_writer.add_image('train flow warped img', tensor2array(flow_ref_img_warped[0]), n_iter) train_writer.add_image( 'train Cam Flow Output', flow_to_image(tensor2array(flow_fwd[0].data[0].cpu())), n_iter) train_writer.add_image( 'train Flow from Depth Output', flow_to_image(tensor2array(flows_cam_fwd.data[0].cpu())), n_iter) train_writer.add_image( 'train Flow and Depth diff', flow_to_image(tensor2array(rigidity_mask_fwd.data[0].cpu())), n_iter) n_iter += 1 return average_loss / i
def train(train_loader, disp_net, pose_net, mask_net, flow_net, optimizer, logger=None, train_writer=None, global_vars_dict=None): # 0. 准备 args = global_vars_dict['args'] n_iter = global_vars_dict['n_iter'] device = global_vars_dict['device'] batch_time = AverageMeter() data_time = AverageMeter() losses = AverageMeter(precision=4) w1, w2, w3, w4 = args.cam_photo_loss_weight, args.mask_loss_weight, args.smooth_loss_weight, args.flow_photo_loss_weight w5 = args.consensus_loss_weight if args.robust: loss_camera = photometric_reconstruction_loss_robust loss_flow = photometric_flow_loss_robust else: loss_camera = photometric_reconstruction_loss loss_flow = photometric_flow_loss #2. switch to train mode disp_net.train() pose_net.train() mask_net.train() flow_net.train() end = time.time() #3. train cycle for i, (tgt_img, ref_imgs, intrinsics, intrinsics_inv) in enumerate(train_loader): # measure data loading time data_time.update(time.time() - end) tgt_img = tgt_img.to(device) ref_imgs = [img.to(device) for img in ref_imgs] intrinsics = intrinsics.to(device) intrinsics_inv = intrinsics_inv.to(device) #3.1 compute output and lossfunc input valve--------------------- #1. disp->depth(none) disparities = disp_net(tgt_img) if args.spatial_normalize: disparities = [spatial_normalize(disp) for disp in disparities] depth = [1 / disp for disp in disparities] #2. pose(none) pose = pose_net(tgt_img, ref_imgs) #pose:[4,4,6] #3.flow_fwd,flow_bwd 全光流 (depth, pose) # 自己改了一点 if args.flownet == 'Back2Future': #临近一共三帧做训练/推断 flow_fwd, flow_bwd, _ = flow_net(tgt_img, ref_imgs[1:3]) elif args.flownet == 'FlowNetC6': flow_fwd = flow_net(tgt_img, ref_imgs[2]) flow_bwd = flow_net(tgt_img, ref_imgs[1]) elif args.flownet == 'FlowNetS': print(' ') # flow_cam 即背景光流 # flow - flow_s = flow_o flow_cam = pose2flow( depth[0].squeeze(), pose[:, 2], intrinsics, intrinsics_inv) # pose[:,2] belongs to forward frame flows_cam_fwd = [ pose2flow(depth_.squeeze(1), pose[:, 2], intrinsics, intrinsics_inv) for depth_ in depth ] flows_cam_bwd = [ pose2flow(depth_.squeeze(1), pose[:, 1], intrinsics, intrinsics_inv) for depth_ in depth ] exp_masks_target = consensus_exp_masks(flows_cam_fwd, flows_cam_bwd, flow_fwd, flow_bwd, tgt_img, ref_imgs[2], ref_imgs[1], wssim=args.wssim, wrig=args.wrig, ws=args.smooth_loss_weight) rigidity_mask_fwd = [ (flows_cam_fwd_i - flow_fwd_i).abs() for flows_cam_fwd_i, flow_fwd_i in zip(flows_cam_fwd, flow_fwd) ] # .normalize() rigidity_mask_bwd = [ (flows_cam_bwd_i - flow_bwd_i).abs() for flows_cam_bwd_i, flow_bwd_i in zip(flows_cam_bwd, flow_bwd) ] # .normalize() #v_u # 4.explainability_mask(none) explainability_mask = mask_net(tgt_img, ref_imgs) #有效区域?4?? #list(5):item:tensor:[4,4,128,512]...[4,4,4,16] value:[0.33~0.48~0.63] #------------------------------------------------- if args.joint_mask_for_depth: explainability_mask_for_depth = compute_joint_mask_for_depth( explainability_mask, rigidity_mask_bwd, rigidity_mask_fwd, args.THRESH) else: explainability_mask_for_depth = explainability_mask #explainability_mask_for_depth list(5) [b,2,h/ , w/] if args.no_non_rigid_mask: flow_exp_mask = [None for exp_mask in explainability_mask] if args.DEBUG: print('Using no masks for flow') else: flow_exp_mask = [ 1 - exp_mask[:, 1:3] for exp_mask in explainability_mask ] # explaninbility mask 本来是背景mask, 背景对应像素为1 #取反改成动物mask,并且只要前后两帧 #list(4) [4,2,256,512] #3.2. compute loss重 # E-r minimizes the photometric loss on static scene if w1 > 0: loss_1 = loss_camera(tgt_img, ref_imgs, intrinsics, intrinsics_inv, depth, explainability_mask_for_depth, pose, lambda_oob=args.lambda_oob, qch=args.qch, wssim=args.wssim) else: loss_1 = torch.tensor([0.]).to(device) # E_M if w2 > 0: loss_2 = explainability_loss( explainability_mask ) #+ 0.2*gaussian_explainability_loss(explainability_mask) else: loss_2 = 0 # E_S if w3 > 0: if args.smoothness_type == "regular": loss_3 = smooth_loss(depth) + smooth_loss( flow_fwd) + smooth_loss(flow_bwd) + smooth_loss( explainability_mask) elif args.smoothness_type == "edgeaware": loss_3 = edge_aware_smoothness_loss( tgt_img, depth) + edge_aware_smoothness_loss( tgt_img, flow_fwd) loss_3 += edge_aware_smoothness_loss( tgt_img, flow_bwd) + edge_aware_smoothness_loss( tgt_img, explainability_mask) else: loss_3 = torch.tensor([0.]).to(device) # E_F # minimizes photometric loss on moving regions if w4 > 0: loss_4 = loss_flow(tgt_img, ref_imgs[1:3], [flow_bwd, flow_fwd], flow_exp_mask, lambda_oob=args.lambda_oob, qch=args.qch, wssim=args.wssim) else: loss_4 = torch.tensor([0.]).to(device) # E_C # drives the collaboration #explainagy_mask:list(6) of [4,4,4,16] rigidity_mask :list(4):[4,2,128,512] if w5 > 0: loss_5 = consensus_depth_flow_mask(explainability_mask, rigidity_mask_bwd, rigidity_mask_fwd, exp_masks_target, exp_masks_target, THRESH=args.THRESH, wbce=args.wbce) else: loss_5 = torch.tensor([0.]).to(device) #3.2.6 loss = w1 * loss_1 + w2 * loss_2 + w3 * loss_3 + w4 * loss_4 + w5 * loss_5 #end of loss #3.3 # record loss and EPE losses.update(loss.item(), args.batch_size) # compute gradient and do Adam step optimizer.zero_grad() loss.backward() optimizer.step() # measure elapsed time batch_time.update(time.time() - end) end = time.time() #3.4 log data # add scalar if args.scalar_freq > 0 and n_iter % args.scalar_freq == 0: train_writer.add_scalar('batch/cam_photometric_error', loss_1.item(), n_iter) if w2 > 0: train_writer.add_scalar('batch/explanability_loss', loss_2.item(), n_iter) train_writer.add_scalar('batch/disparity_smoothness_loss', loss_3.item(), n_iter) train_writer.add_scalar('batch/flow_photometric_error', loss_4.item(), n_iter) train_writer.add_scalar('batch/consensus_error', loss_5.item(), n_iter) train_writer.add_scalar('batch/total_loss', loss.item(), n_iter) # add_image为0 则不输出 if args.training_output_freq > 0 and n_iter % args.training_output_freq == 0: train_writer.add_image('train Input', tensor2array(tgt_img[0]), n_iter) train_writer.add_image( 'train Cam Flow Output', flow_to_image(tensor2array(flow_cam.data[0].cpu())), n_iter) for k, scaled_depth in enumerate(depth): train_writer.add_image( 'train Dispnet Output Normalized111 {}'.format(k), tensor2array(disparities[k].data[0].cpu(), max_value=None, colormap='bone'), n_iter) train_writer.add_image( 'train Depth Output {}'.format(k), tensor2array(1 / disparities[k].data[0].cpu(), max_value=10), n_iter) train_writer.add_image( 'train Non Rigid Flow Output {}'.format(k), flow_to_image(tensor2array(flow_fwd[k].data[0].cpu())), n_iter) train_writer.add_image( 'train Target Rigidity {}'.format(k), tensor2array((rigidity_mask_fwd[k] > args.THRESH).type_as( rigidity_mask_fwd[k]).data[0].cpu(), max_value=1, colormap='bone'), n_iter) b, _, h, w = scaled_depth.size() downscale = tgt_img.size(2) / h tgt_img_scaled = nn.functional.adaptive_avg_pool2d( tgt_img, (h, w)) ref_imgs_scaled = [ nn.functional.adaptive_avg_pool2d(ref_img, (h, w)) for ref_img in ref_imgs ] intrinsics_scaled = torch.cat( (intrinsics[:, 0:2] / downscale, intrinsics[:, 2:]), dim=1) intrinsics_scaled_inv = torch.cat( (intrinsics_inv[:, :, 0:2] * downscale, intrinsics_inv[:, :, 2:]), dim=2) train_writer.add_image( 'train Non Rigid Warped Image {}'.format(k), tensor2array( flow_warp(ref_imgs_scaled[2], flow_fwd[k]).data[0].cpu()), n_iter) # log warped images along with explainability mask for j, ref in enumerate(ref_imgs_scaled): ref_warped = inverse_warp( ref, scaled_depth[:, 0], pose[:, j], intrinsics_scaled, intrinsics_scaled_inv, rotation_mode=args.rotation_mode, padding_mode=args.padding_mode)[0] train_writer.add_image( 'train Warped Outputs {} {}'.format(k, j), tensor2array(ref_warped.data.cpu()), n_iter) train_writer.add_image( 'train Diff Outputs {} {}'.format(k, j), tensor2array( 0.5 * (tgt_img_scaled[0] - ref_warped).abs().data.cpu()), n_iter) if explainability_mask[k] is not None: train_writer.add_image( 'train Exp mask Outputs {} {}'.format(k, j), tensor2array(explainability_mask[k][0, j].data.cpu(), max_value=1, colormap='bone'), n_iter) # csv file write with open(args.save_path / args.log_full, 'a') as csvfile: writer = csv.writer(csvfile, delimiter='\t') writer.writerow([ loss.item(), loss_1.item(), loss_2.item() if w2 > 0 else 0, loss_3.item(), loss_4.item() ]) #terminal output if args.log_terminal: logger.train_bar.update(i + 1) #当前epoch 进度 if i % args.print_freq == 0: logger.valid_bar_writer.write( 'Train: Time {} Data {} Loss {}'.format( batch_time, data_time, losses)) # 3.4 edge conditionsssssssssssssssssssssssss epoch_size = len(train_loader) if i >= epoch_size - 1: break n_iter += 1 global_vars_dict['n_iter'] = n_iter return losses.avg[0] #epoch loss
def main(): global args args = parser.parse_args() args.pretrained_path = Path(args.pretrained_path) if args.output_dir is not None: args.output_dir = Path(args.output_dir) args.output_dir.makedirs_p() image_dir = args.output_dir / 'images' mask_dir = args.output_dir / 'mask' viz_dir = args.output_dir / 'viz' testing_dir = args.output_dir / 'testing' testing_dir_flo = args.output_dir / 'testing_flo' image_dir.makedirs_p() mask_dir.makedirs_p() viz_dir.makedirs_p() testing_dir.makedirs_p() testing_dir_flo.makedirs_p() normalize = custom_transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) flow_loader_h, flow_loader_w = 256, 832 valid_flow_transform = custom_transforms.Compose([ custom_transforms.Scale(h=flow_loader_h, w=flow_loader_w), custom_transforms.ArrayToTensor(), normalize ]) val_flow_set = KITTI2015Test(root=args.kitti_dir, sequence_length=5, transform=valid_flow_transform) if args.DEBUG: print("DEBUG MODE: Using Training Set") val_flow_set = KITTI2015Test(root=args.kitti_dir, sequence_length=5, transform=valid_flow_transform, phase='training') val_loader = torch.utils.data.DataLoader(val_flow_set, batch_size=1, shuffle=False, num_workers=2, pin_memory=True, drop_last=True) disp_net = getattr(models, args.dispnet)().cuda() pose_net = getattr(models, args.posenet)(nb_ref_imgs=4).cuda() mask_net = getattr(models, args.masknet)(nb_ref_imgs=4).cuda() flow_net = getattr(models, args.flownet)(nlevels=args.nlevels).cuda() dispnet_weights = torch.load(args.pretrained_path / 'dispnet_model_best.pth.tar') posenet_weights = torch.load(args.pretrained_path / 'posenet_model_best.pth.tar') masknet_weights = torch.load(args.pretrained_path / 'masknet_model_best.pth.tar') flownet_weights = torch.load(args.pretrained_path / 'flownet_model_best.pth.tar') disp_net.load_state_dict(dispnet_weights['state_dict']) pose_net.load_state_dict(posenet_weights['state_dict']) flow_net.load_state_dict(flownet_weights['state_dict']) mask_net.load_state_dict(masknet_weights['state_dict']) disp_net.eval() pose_net.eval() mask_net.eval() flow_net.eval() for i, (tgt_img, ref_imgs, intrinsics, intrinsics_inv, tgt_img_original) in enumerate(tqdm(val_loader)): tgt_img_var = Variable(tgt_img.cuda(), volatile=True) ref_imgs_var = [ Variable(img.cuda(), volatile=True) for img in ref_imgs ] intrinsics_var = Variable(intrinsics.cuda(), volatile=True) intrinsics_inv_var = Variable(intrinsics_inv.cuda(), volatile=True) disp = disp_net(tgt_img_var) depth = 1 / disp pose = pose_net(tgt_img_var, ref_imgs_var) explainability_mask = mask_net(tgt_img_var, ref_imgs_var) if args.flownet == 'Back2Future': flow_fwd, _, _ = flow_net(tgt_img_var, ref_imgs_var[1:3]) else: flow_fwd = flow_net(tgt_img_var, ref_imgs_var[2]) flow_cam = pose2flow(depth.squeeze(1), pose[:, 2], intrinsics_var, intrinsics_inv_var) rigidity_mask = 1 - (1 - explainability_mask[:, 1]) * ( 1 - explainability_mask[:, 2]).unsqueeze(1) > 0.5 rigidity_mask_census_soft = (flow_cam - flow_fwd).abs() #.normalize() rigidity_mask_census_u = rigidity_mask_census_soft[:, 0] < args.THRESH rigidity_mask_census_v = rigidity_mask_census_soft[:, 1] < args.THRESH rigidity_mask_census = (rigidity_mask_census_u).type_as(flow_fwd) * ( rigidity_mask_census_v).type_as(flow_fwd) rigidity_mask_combined = 1 - ( 1 - rigidity_mask.type_as(explainability_mask)) * ( 1 - rigidity_mask_census.type_as(explainability_mask)) _, _, h_pred, w_pred = flow_cam.size() _, _, h_gt, w_gt = tgt_img_original.size() rigidity_pred_mask = nn.functional.upsample(rigidity_mask_combined, size=(h_pred, w_pred), mode='bilinear') non_rigid_pred = (rigidity_pred_mask <= args.THRESH ).type_as(flow_fwd).expand_as(flow_fwd) * flow_fwd rigid_pred = (rigidity_pred_mask > args.THRESH ).type_as(flow_cam).expand_as(flow_cam) * flow_cam total_pred = non_rigid_pred + rigid_pred pred_fullres = nn.functional.upsample(total_pred, size=(h_gt, w_gt), mode='bilinear') pred_fullres[:, 0, :, :] = pred_fullres[:, 0, :, :] * (w_gt / w_pred) pred_fullres[:, 1, :, :] = pred_fullres[:, 1, :, :] * (h_gt / h_pred) flow_fwd_fullres = nn.functional.upsample(flow_fwd, size=(h_gt, w_gt), mode='bilinear') flow_fwd_fullres[:, 0, :, :] = flow_fwd_fullres[:, 0, :, :] * (w_gt / w_pred) flow_fwd_fullres[:, 1, :, :] = flow_fwd_fullres[:, 1, :, :] * (h_gt / h_pred) flow_cam_fullres = nn.functional.upsample(flow_cam, size=(h_gt, w_gt), mode='bilinear') flow_cam_fullres[:, 0, :, :] = flow_cam_fullres[:, 0, :, :] * (w_gt / w_pred) flow_cam_fullres[:, 1, :, :] = flow_cam_fullres[:, 1, :, :] * (h_gt / h_pred) tgt_img_np = tgt_img[0].numpy() rigidity_mask_combined_np = rigidity_mask_combined.cpu().data[0].numpy( ) if args.output_dir is not None: np.save(image_dir / str(i).zfill(3), tgt_img_np) np.save(mask_dir / str(i).zfill(3), rigidity_mask_combined_np) pred_u = pred_fullres[0][0].data.cpu().numpy() pred_v = pred_fullres[0][1].data.cpu().numpy() flow_io.flow_write_png(testing_dir / str(i).zfill(6) + '_10.png', u=pred_u, v=pred_v) flow_io.flow_write(testing_dir_flo / str(i).zfill(6) + '_10.flo', pred_u, pred_v) if (args.output_dir is not None): ind = int(i) tgt_img_viz = tensor2array(tgt_img[0].cpu()) depth_viz = tensor2array(disp.data[0].cpu(), max_value=None, colormap='magma') mask_viz = tensor2array(rigidity_mask_combined.data[0].cpu(), max_value=1, colormap='magma') row2_viz = flow_to_image( np.hstack((tensor2array(flow_cam_fullres.data[0].cpu()), tensor2array(flow_fwd_fullres.data[0].cpu()), tensor2array(pred_fullres.data[0].cpu())))) row1_viz = np.hstack((tgt_img_viz, depth_viz, mask_viz)) row1_viz_im = Image.fromarray( (255 * row1_viz.transpose(1, 2, 0)).astype('uint8')) row2_viz_im = Image.fromarray( (255 * row2_viz.transpose(1, 2, 0)).astype('uint8')) row1_viz_im.save(viz_dir / str(i).zfill(3) + '01.png') row2_viz_im.save(viz_dir / str(i).zfill(3) + '02.png') print("Done!")
def validate_flow_with_gt(val_loader, disp_net, pose_net, mask_net, flow_net, epoch, logger, output_writers=[]): global args batch_time = AverageMeter() error_names = [ 'epe_total', 'epe_rigid', 'epe_non_rigid', 'outliers', 'epe_total_with_gt_mask', 'epe_rigid_with_gt_mask', 'epe_non_rigid_with_gt_mask', 'outliers_gt_mask' ] errors = AverageMeter(i=len(error_names)) log_outputs = len(output_writers) > 0 # switch to evaluate mode disp_net.eval() pose_net.eval() mask_net.eval() flow_net.eval() end = time.time() poses = np.zeros( ((len(val_loader) - 1) * 1 * (args.sequence_length - 1), 6)) for i, (tgt_img, ref_imgs, intrinsics, intrinsics_inv, flow_gt, obj_map_gt) in enumerate(val_loader): tgt_img = Variable(tgt_img.cuda(), volatile=True) ref_imgs = [Variable(img.cuda(), volatile=True) for img in ref_imgs] intrinsics_var = Variable(intrinsics.cuda(), volatile=True) intrinsics_inv_var = Variable(intrinsics_inv.cuda(), volatile=True) flow_gt_var = Variable(flow_gt.cuda(), volatile=True) obj_map_gt_var = Variable(obj_map_gt.cuda(), volatile=True) # compute output------------------------- #1. disp fwd disp = disp_net(tgt_img) if args.spatial_normalize: disp = spatial_normalize(disp) depth = 1 / disp #2. pose fwd pose = pose_net(tgt_img, ref_imgs) #3. mask fwd explainability_mask = mask_net(tgt_img, ref_imgs) #4. flow fwd if args.flownet == 'Back2Future': flow_fwd, flow_bwd, _ = flow_net(tgt_img, ref_imgs[1:3]) #前一帧,后一阵 elif args.flownet == 'FlowNetC6': flow_fwd = flow_net(tgt_img, ref_imgs[2]) flow_bwd = flow_net(tgt_img, ref_imgs[1]) # compute output------------------------- if args.DEBUG: flow_fwd_x = flow_fwd[:, 0].view(-1).abs().data print("Flow Fwd Median: ", flow_fwd_x.median()) flow_gt_var_x = flow_gt_var[:, 0].view(-1).abs().data print( "Flow GT Median: ", flow_gt_var_x.index_select( 0, flow_gt_var_x.nonzero().view(-1)).median()) flow_cam = pose2flow(depth.squeeze(1), pose[:, 2], intrinsics_var, intrinsics_inv_var) oob_rigid = flow2oob(flow_cam) oob_non_rigid = flow2oob(flow_fwd) rigidity_mask = 1 - (1 - explainability_mask[:, 1]) * ( 1 - explainability_mask[:, 2]).unsqueeze(1) > 0.5 rigidity_mask_census_soft = (flow_cam - flow_fwd).abs() #.normalize() rigidity_mask_census_u = rigidity_mask_census_soft[:, 0] < args.THRESH rigidity_mask_census_v = rigidity_mask_census_soft[:, 1] < args.THRESH rigidity_mask_census = (rigidity_mask_census_u).type_as(flow_fwd) * ( rigidity_mask_census_v).type_as(flow_fwd) rigidity_mask_combined = 1 - ( 1 - rigidity_mask.type_as(explainability_mask)) * ( 1 - rigidity_mask_census.type_as(explainability_mask)) #get flow flow_fwd_non_rigid = (rigidity_mask_combined <= args.THRESH).type_as( flow_fwd).expand_as(flow_fwd) * flow_fwd flow_fwd_rigid = (rigidity_mask_combined > args.THRESH ).type_as(flow_fwd).expand_as(flow_fwd) * flow_cam total_flow = flow_fwd_rigid + flow_fwd_non_rigid obj_map_gt_var_expanded = obj_map_gt_var.unsqueeze(1).type_as(flow_fwd) if log_outputs and i % 10 == 0 and i / 10 < len(output_writers): index = int(i // 10) if epoch == 0: output_writers[index].add_image('val flow Input', tensor2array(tgt_img[0]), 0) flow_to_show = flow_gt[0][:2, :, :].cpu() output_writers[index].add_image( 'val target Flow', flow_to_image(tensor2array(flow_to_show)), epoch) output_writers[index].add_image( 'val Total Flow Output', flow_to_image(tensor2array(total_flow.data[0].cpu())), epoch) output_writers[index].add_image( 'val Rigid Flow Output', flow_to_image(tensor2array(flow_fwd_rigid.data[0].cpu())), epoch) output_writers[index].add_image( 'val Non-rigid Flow Output', flow_to_image(tensor2array(flow_fwd_non_rigid.data[0].cpu())), epoch) output_writers[index].add_image( 'val Out of Bound (Rigid)', tensor2array(oob_rigid.type(torch.FloatTensor).data[0].cpu(), max_value=1, colormap='bone'), epoch) output_writers[index].add_scalar( 'val Mean oob (Rigid)', oob_rigid.type(torch.FloatTensor).sum(), epoch) output_writers[index].add_image( 'val Out of Bound (Non-Rigid)', tensor2array(oob_non_rigid.type( torch.FloatTensor).data[0].cpu(), max_value=1, colormap='bone'), epoch) output_writers[index].add_scalar( 'val Mean oob (Non-Rigid)', oob_non_rigid.type(torch.FloatTensor).sum(), epoch) output_writers[index].add_image( 'val Cam Flow Errors', tensor2array(flow_diff(flow_gt_var, flow_cam).data[0].cpu()), epoch) output_writers[index].add_image( 'val Rigidity Mask', tensor2array(rigidity_mask.data[0].cpu(), max_value=1, colormap='bone'), epoch) output_writers[index].add_image( 'val Rigidity Mask Census', tensor2array(rigidity_mask_census.data[0].cpu(), max_value=1, colormap='bone'), epoch) for j, ref in enumerate(ref_imgs): ref_warped = inverse_warp(ref[:1], depth[:1, 0], pose[:1, j], intrinsics_var[:1], intrinsics_inv_var[:1], rotation_mode=args.rotation_mode, padding_mode=args.padding_mode)[0] output_writers[index].add_image( 'val Warped Outputs {}'.format(j), tensor2array(ref_warped.data.cpu()), epoch) output_writers[index].add_image( 'val Diff Outputs {}'.format(j), tensor2array(0.5 * (tgt_img[0] - ref_warped).abs().data.cpu()), epoch) if explainability_mask is not None: output_writers[index].add_image( 'val Exp mask Outputs {}'.format(j), tensor2array(explainability_mask[0, j].data.cpu(), max_value=1, colormap='bone'), epoch) if args.DEBUG: # Check if pose2flow is consistant with inverse warp ref_warped_from_depth = inverse_warp( ref_imgs[2][:1], depth[:1, 0], pose[:1, 2], intrinsics_var[:1], intrinsics_inv_var[:1], rotation_mode=args.rotation_mode, padding_mode=args.padding_mode)[0] ref_warped_from_cam_flow = flow_warp(ref_imgs[2][:1], flow_cam)[0] print( "DEBUG_INFO: Inverse_warp vs pose2flow", torch.mean( torch.abs(ref_warped_from_depth - ref_warped_from_cam_flow)).item()) output_writers[index].add_image( 'val Warped Outputs from Cam Flow', tensor2array(ref_warped_from_cam_flow.data.cpu()), epoch) output_writers[index].add_image( 'val Warped Outputs from inverse warp', tensor2array(ref_warped_from_depth.data.cpu()), epoch) if log_outputs and i < len(val_loader) - 1: step = args.sequence_length - 1 poses[i * step:(i + 1) * step] = pose.data.cpu().view(-1, 6).numpy() if np.isnan(flow_gt.sum().item()) or np.isnan( total_flow.data.sum().item()): print('NaN encountered') # _epe_errors = compute_all_epes( flow_gt_var, flow_cam, flow_fwd, rigidity_mask_combined) + compute_all_epes( flow_gt_var, flow_cam, flow_fwd, (1 - obj_map_gt_var_expanded)) errors.update(_epe_errors) if args.DEBUG: print("DEBUG_INFO: EPE errors: ", _epe_errors) # measure elapsed time batch_time.update(time.time() - end) end = time.time() if log_outputs: output_writers[0].add_histogram('val poses_tx', poses[:, 0], epoch) output_writers[0].add_histogram('val poses_ty', poses[:, 1], epoch) output_writers[0].add_histogram('val poses_tz', poses[:, 2], epoch) if args.rotation_mode == 'euler': rot_coeffs = ['rx', 'ry', 'rz'] elif args.rotation_mode == 'quat': rot_coeffs = ['qx', 'qy', 'qz'] output_writers[0].add_histogram('val poses_{}'.format(rot_coeffs[0]), poses[:, 3], epoch) output_writers[0].add_histogram('val poses_{}'.format(rot_coeffs[1]), poses[:, 4], epoch) output_writers[0].add_histogram('val poses_{}'.format(rot_coeffs[2]), poses[:, 5], epoch) if args.DEBUG: print("DEBUG_INFO =================>") print("DEBUG_INFO: Average EPE : ", errors.avg) print("DEBUG_INFO =================>") print("DEBUG_INFO =================>") print("DEBUG_INFO =================>") return errors.avg, error_names
def train(train_loader, flow_net, optimizer, epoch_size, logger=None, train_writer=None): global args, n_iter batch_time = AverageMeter() data_time = AverageMeter() losses = AverageMeter(precision=4) # switch to train mode flow_net.train() end = time.time() for i, (tgt_img, ref_imgs, intrinsics, intrinsics_inv) in enumerate(train_loader): # measure data loading time data_time.update(time.time() - end) tgt_img_var = Variable(tgt_img.cuda()) ref_imgs_var = [Variable(img.cuda()) for img in ref_imgs] if args.flownet == 'Back2Future': flow_fwd, flow_bwd = flow_net(tgt_img_var, ref_imgs_var) else: flow_fwd = flow_net(tgt_img_var, ref_imgs_var[1]) flow_bwd = flow_net(tgt_img_var, ref_imgs_var[0]) loss_smooth = torch.zeros(1).cuda() loss_flow_recon = torch.zeros(1).cuda() loss_velocity_consis = torch.zeros(1).cuda() if args.flow_photo_loss_weight_first: if args.min: loss_flow_recon += args.flow_photo_loss_weight_first*photometric_flow_min_loss(tgt_img_var, ref_imgs_var, [flow_bwd, flow_fwd], lambda_oob=args.lambda_oob, qch=args.qch, wssim=args.wssim) else: loss_flow_recon += args.flow_photo_loss_weight_first*photometric_flow_loss(tgt_img_var, ref_imgs_var, [flow_bwd, flow_fwd], lambda_oob=args.lambda_oob, qch=args.qch, wssim=args.wssim) if args.flow_photo_loss_weight_second: if args.min: loss_per, loss_weight= photometric_flow_gradient_min_loss(tgt_img_var, ref_imgs_var, [flow_bwd, flow_fwd], lambda_oob=args.lambda_oob, qch=args.qch, wssim=args.wssim) loss_flow_recon += args.flow_photo_loss_weight_second * loss_per else: loss_flow_recon += args.flow_photo_loss_weight_second*photometric_flow_gradient_loss(tgt_img_var, ref_imgs_var, [flow_bwd, flow_fwd], lambda_oob=args.lambda_oob, qch=args.qch, wssim=args.wssim) if args.smooth_loss_weight_first: if args.smoothness_type == "regular": loss_smooth += args.smooth_loss_weight_first*(smooth_loss(flow_fwd) + smooth_loss(flow_bwd)) elif args.smoothness_type == "edgeaware": loss_smooth += args.smooth_loss_weight_first*(edge_aware_smoothness_loss(tgt_img_var, flow_fwd)+edge_aware_smoothness_loss(tgt_img_var, flow_bwd)) if args.smooth_loss_weight_second: if args.smoothness_type == "regular": loss_smooth += args.smooth_loss_weight_second*(smooth_loss(flow_fwd) + smooth_loss(flow_bwd)) elif args.smoothness_type == "edgeaware": loss_smooth = args.smooth_loss_weight_second*(edge_aware_smoothness_second_order_loss_change_weight(tgt_img_var, flow_bwd, args.alpha)\ + edge_aware_smoothness_second_order_loss_change_weight(tgt_img_var, flow_fwd, args.alpha)) if args.velocity_consis_loss_weight: loss_velocity_consis = args.velocity_consis_loss_weight*flow_velocity_consis_loss( [flow_bwd, flow_fwd]) loss = loss_smooth + loss_flow_recon + loss_velocity_consis if i > 0 and n_iter % args.print_freq == 0: train_writer.add_scalar('flow_photometric_error', loss_flow_recon.item(), n_iter) train_writer.add_scalar('flow_smoothness_loss', loss_smooth.item(), n_iter) train_writer.add_scalar('velocity_consis_loss', loss_velocity_consis.item(), n_iter) train_writer.add_scalar('total_loss', loss.item(), n_iter) if args.training_output_freq > 0 and n_iter % args.training_output_freq == 0: train_writer.add_image('train Input', tensor2array(tgt_img[0]), n_iter) train_writer.add_image('train Flow FWD Output',flow_to_image(tensor2array(flow_fwd[0].data[0].cpu())) , n_iter ) train_writer.add_image('train Flow BWD Output',flow_to_image(tensor2array(flow_bwd[0].data[0].cpu())) , n_iter ) loss_weight_bwd = loss_weight[0][0,0,:,:].unsqueeze(0) loss_weight_fwd = loss_weight[0][0,1,:,:].unsqueeze(0) train_writer.add_image('loss_weight_bwd', tensor2array(loss_weight_bwd.data[0].cpu(), max_value=None, colormap='bone'), n_iter) train_writer.add_image('loss_weight_fwd', tensor2array(loss_weight_fwd.data[0].cpu(), max_value=None, colormap='bone'), n_iter) train_writer.add_image('train Flow FWD error Image',tensor2array(flow_warp(tgt_img_var-ref_imgs_var[1],flow_fwd[0]).data[0].cpu()) , n_iter ) train_writer.add_image('train Flow BWD error Image',tensor2array(flow_warp(tgt_img_var-ref_imgs_var[0],flow_bwd[0]).data[0].cpu()) , n_iter ) # record loss and EPE losses.update(loss.item(), args.batch_size) # compute gradient and do Adam step optimizer.zero_grad() loss.backward() optimizer.step() # measure elapsed time batch_time.update(time.time() - end) end = time.time() if args.log_terminal: logger.train_bar.update(i+1) if i % args.print_freq == 0: logger.train_writer.write('Train: Time {} Data {} Loss {}'.format(batch_time, data_time, losses)) if i >= epoch_size - 1: break n_iter += 1 return losses.avg[0]
def main(): global args args = parser.parse_args() save_path = Path(args.name) args.save_path = 'results' / save_path #/timestamp print('=> will save everything to {}'.format(args.save_path)) args.save_path.makedirs_p() output_vis_dir = args.save_path / 'images' output_vis_dir.makedirs_p() args.batch_size = 1 output_writer = SummaryWriter(args.save_path / 'valid') # Data loading code flow_loader_h, flow_loader_w = 384, 1280 normalize = custom_transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) # valid_transform = custom_transforms.Compose([custom_transforms.Scale(h=flow_loader_h, w=flow_loader_w), # custom_transforms.ArrayToTensor(), normalize]) valid_transform = custom_transforms.Compose([ custom_transforms.Scale(h=flow_loader_h, w=flow_loader_w), custom_transforms.ArrayToTensor() ]) if args.valset == "kitti2015": # from datasets.validation_flow import ValidationFlowKitti2015MV # val_set = ValidationFlowKitti2015MV(root='/ps/project/datasets/AllFlowData/kitti/kitti2015', transform=valid_transform, compression=args.compression, raw_root='/is/rg/avg/jjanai/data/Kitti_2012_2015/Raw', example=args.example, true_motion=args.true_motion) from datasets.validation_flow import ValidationFlowKitti2015 # # val_set = ValidationFlowKitti2015(root='/is/ps2/aranjan/AllFlowData/kitti/kitti2015', transform=valid_transform, compression=args.compression) val_set = ValidationFlowKitti2015( root='/ps/project/datasets/AllFlowData/kitti/kitti2015', transform=valid_transform, compression=args.compression, raw_root='/is/rg/avg/jjanai/data/Kitti_2012_2015/Raw', example=args.example, true_motion=args.true_motion) elif args.valset == "kitti2012": from datasets.validation_flow import ValidationFlowKitti2012 # val_set = ValidationFlowKitti2012(root='/is/ps2/aranjan/AllFlowData/kitti/kitti2012', transform=valid_transform, compression=args.compression) val_set = ValidationFlowKitti2012( root='/ps/project/datasets/AllFlowData/kitti/kitti2012', transform=valid_transform, compression=args.compression, raw_root='/is/rg/avg/jjanai/data/Kitti_2012_2015/Raw') print('{} samples found in valid scenes'.format(len(val_set))) val_loader = torch.utils.data.DataLoader( val_set, batch_size= 1, # batch size is 1 since images in kitti have different sizes shuffle=False, num_workers=args.workers, pin_memory=True, drop_last=True) result_file = open(os.path.join(args.save_path, 'results.csv'), 'a') result_scene_file = open(os.path.join(args.save_path, 'result_scenes.csv'), 'a') # create model print("=> fetching model") if args.flownet == 'SpyNet': flow_net = getattr(models, args.flownet)(nlevels=6, pretrained=True) elif args.flownet == 'Back2Future': flow_net = getattr( models, args.flownet)(pretrained='pretrained/b2f_rm_hard.pth.tar') elif args.flownet == 'PWCNet': flow_net = models.pwc_dc_net( 'pretrained/pwc_net_chairs.pth.tar') # pwc_net.pth.tar') else: flow_net = getattr(models, args.flownet)() if args.flownet in ['SpyNet', 'Back2Future', 'PWCNet']: print("=> using pre-trained weights for " + args.flownet) elif args.flownet in ['FlowNetC']: print("=> using pre-trained weights for FlowNetC") weights = torch.load('pretrained/FlowNet2-C_checkpoint.pth.tar') flow_net.load_state_dict(weights['state_dict']) elif args.flownet in ['FlowNetS']: print("=> using pre-trained weights for FlowNetS") weights = torch.load('pretrained/flownets.pth.tar') flow_net.load_state_dict(weights['state_dict']) elif args.flownet in ['FlowNet2']: print("=> using pre-trained weights for FlowNet2") weights = torch.load('pretrained/FlowNet2_checkpoint.pth.tar') flow_net.load_state_dict(weights['state_dict']) else: flow_net.init_weights() flow_net = flow_net.cuda() cudnn.benchmark = True if args.whole_img == 0 and args.compression == 0: print("Loading patch from ", args.patch_path) patch = torch.load(args.patch_path) patch_shape = patch.shape if args.mask_path: mask_image = load_as_float(args.mask_path) mask_image = imresize(mask_image, (patch_shape[-1], patch_shape[-2])) / 256. mask = np.array([mask_image.transpose(2, 0, 1)]) else: if args.patch_type == 'circle': mask = createCircularMask(patch_shape[-2], patch_shape[-1]).astype('float32') mask = np.array([[mask, mask, mask]]) elif args.patch_type == 'square': mask = np.ones(patch_shape) else: # add gaussian noise mean = 0 var = 1 sigma = var**0.5 patch = np.random.normal(mean, sigma, (flow_loader_h, flow_loader_w, 3)) patch = patch.reshape(3, flow_loader_h, flow_loader_w) mask = np.ones(patch.shape) * args.whole_img #import ipdb; ipdb.set_trace() error_names = ['epe', 'adv_epe', 'cos_sim', 'adv_cos_sim'] errors = AverageMeter(i=len(error_names)) # header result_file.write("{:>10}, {:>10}, {:>10}, {:>10}\n".format(*error_names)) result_scene_file.write("{:>10}, {:>10}, {:>10}, {:>10}, {:>10}\n".format( *(['scene'] + error_names))) flow_net.eval() # set seed for reproductivity np.random.seed(1337) for i, (ref_img_past, tgt_img, ref_img, flow_gt, disp_gt, calib, poses) in enumerate(tqdm(val_loader)): tgt_img_var = Variable(tgt_img.cuda(), volatile=True) ref_past_img_var = Variable(ref_img_past.cuda(), volatile=True) ref_img_var = Variable(ref_img.cuda(), volatile=True) flow_gt_var = Variable(flow_gt.cuda(), volatile=True) if type(flow_net).__name__ == 'Back2Future': flow_fwd = flow_net(ref_past_img_var, tgt_img_var, ref_img_var) else: flow_fwd = flow_net(tgt_img_var, ref_img_var) data_shape = tgt_img.cpu().numpy().shape margin = 0 if len(calib) > 0: margin = int(disp_gt.max()) random_x = args.fixed_loc_x random_y = args.fixed_loc_y if args.whole_img == 0: if args.patch_type == 'circle': patch_full, mask_full, _, random_x, random_y, _ = circle_transform( patch, mask, patch.copy(), data_shape, patch_shape, margin, norotate=args.norotate, fixed_loc=(random_x, random_y)) elif args.patch_type == 'square': patch_full, mask_full, _, _, _ = square_transform( patch, mask, patch.copy(), data_shape, patch_shape, norotate=args.norotate) patch_full, mask_full = torch.FloatTensor( patch_full), torch.FloatTensor(mask_full) else: patch_full, mask_full = torch.FloatTensor( patch), torch.FloatTensor(mask) patch_full, mask_full = patch_full.cuda(), mask_full.cuda() patch_var, mask_var = Variable(patch_full), Variable(mask_full) patch_var_future = patch_var_past = patch_var mask_var_future = mask_var_past = mask_var # adverserial flow bt, _, h_gt, w_gt = flow_gt_var.shape forward_patch_flow = Variable(torch.cat((torch.zeros( (bt, 2, h_gt, w_gt)), torch.ones((bt, 1, h_gt, w_gt))), 1).cuda(), volatile=True) # project patch into 3D scene if len(calib) > 0: # #################################### ONLY WORKS WITH BATCH SIZE 1 #################################### imu2vel = calib['imu2vel']["RT"][0].numpy() imu2cam = calib['P_imu_cam'][0].numpy() imu2img = calib['P_imu_img'][0].numpy() pose_past = poses[0][0].numpy() pose_ref = poses[1][0].numpy() inv_pose_ref = inv(pose_ref) pose_fut = poses[2][0].numpy() # get point in IMU patch_disp = disp_gt[0, random_y:random_y + patch_shape[-2], random_x:random_x + patch_shape[-1]] valid = (patch_disp > 0) # set to object or free space disparity if False and args.fixed_loc_x > 0 and args.fixed_loc_y > 0: # disparity = patch_disp[valid].mean() - 3 # small correction for gps errors disparity = patch_disp[valid].mean() else: subset = patch_disp[valid] min_disp = 0 if len(subset) > 0: min_disp = subset.min() max_disp = disp_gt.max() disparity = np.random.uniform(min_disp, max_disp) # disparity # print('Disp from ', min_disp, ' to ', max_disp) depth = (calib['cam']['focal_length_x'] * calib['cam']['baseline'] / disparity) p_cam0 = np.array([[0], [0], [0], [1]]) p_cam0[0] = depth * ( random_x - calib['cam']['cx']) / calib['cam']['focal_length_x'] p_cam0[1] = depth * ( random_y - calib['cam']['cy']) / calib['cam']['focal_length_y'] p_cam0[2] = depth # transform T_p_cam0 = np.eye(4) T_p_cam0[0:4, 3:4] = p_cam0 # transformation to generate patch points patch_size = -0.25 pts = np.array([[0, 0, 0, 1], [0, patch_size, 0, 1], [patch_size, 0, 0, 1], [patch_size, patch_size, 0, 1]]).T pts = inv(imu2cam).dot(T_p_cam0.dot(pts)) # get points in reference image pts_src = pose_ref.dot(pts) pts_src = imu2img.dot(pts_src) pts_src = pts_src[:3, :] / pts_src[2:3, :].repeat(3, 0) # get points in past image pts_past = pose_past.dot(pts) pts_past = imu2img.dot(pts_past) pts_past = pts_past[:3, :] / pts_past[2:3, :].repeat(3, 0) # get points in future image pts_fut = pose_fut.dot(pts) pts_fut = imu2img.dot(pts_fut) pts_fut = pts_fut[:3, :] / pts_fut[2:3, :].repeat(3, 0) # find homography between points H_past, _ = cv2.findHomography(pts_src.T, pts_past.T, cv2.RANSAC) H_fut, _ = cv2.findHomography(pts_src.T, pts_fut.T, cv2.RANSAC) # import pdb; pdb.set_trace() refMtrx = torch.from_numpy(H_fut).float().cuda() refMtrx = refMtrx.repeat(args.batch_size, 1, 1) # get pixel origins X, Y = np.meshgrid(np.arange(flow_loader_w), np.arange(flow_loader_h)) X, Y = X.flatten(), Y.flatten() XYhom = np.stack([X, Y, np.ones_like(X)], axis=1).T XYhom = np.tile(XYhom, [args.batch_size, 1, 1]).astype(np.float32) XYhom = torch.from_numpy(XYhom).cuda() XHom, YHom, Zom = torch.unbind(XYhom, dim=1) XHom = XHom.resize_( (args.batch_size, flow_loader_h, flow_loader_w)) YHom = YHom.resize_( (args.batch_size, flow_loader_h, flow_loader_w)) # warp the canonical coordinates XYwarpHom = refMtrx.matmul(XYhom) XwarpHom, YwarpHom, ZwarpHom = torch.unbind(XYwarpHom, dim=1) Xwarp = (XwarpHom / (ZwarpHom + 1e-8)).resize_( (args.batch_size, flow_loader_h, flow_loader_w)) Ywarp = (YwarpHom / (ZwarpHom + 1e-8)).resize_( (args.batch_size, flow_loader_h, flow_loader_w)) # get forward flow u = (XHom - Xwarp).unsqueeze(1) v = (YHom - Ywarp).unsqueeze(1) flow = torch.cat((u, v), 1) flow = nn.functional.upsample(flow, size=(h_gt, w_gt), mode='bilinear') flow[:, 0, :, :] = flow[:, 0, :, :] * (w_gt / flow_loader_w) flow[:, 1, :, :] = flow[:, 1, :, :] * (h_gt / flow_loader_h) forward_patch_flow[:, :2, :, :] = flow # get grid for resampling Xwarp = 2 * ((Xwarp / (flow_loader_w - 1)) - 0.5) Ywarp = 2 * ((Ywarp / (flow_loader_h - 1)) - 0.5) grid = torch.stack([Xwarp, Ywarp], dim=-1) # sampling with bilinear interpolation patch_var_future = torch.nn.functional.grid_sample(patch_var, grid, mode="bilinear") mask_var_future = torch.nn.functional.grid_sample(mask_var, grid, mode="bilinear") # use past homography refMtrxP = torch.from_numpy(H_past).float().cuda() refMtrx = refMtrx.repeat(args.batch_size, 1, 1) # warp the canonical coordinates XYwarpHomP = refMtrxP.matmul(XYhom) XwarpHomP, YwarpHomP, ZwarpHomP = torch.unbind(XYwarpHomP, dim=1) XwarpP = (XwarpHomP / (ZwarpHomP + 1e-8)).resize_( (args.batch_size, flow_loader_h, flow_loader_w)) YwarpP = (YwarpHomP / (ZwarpHomP + 1e-8)).resize_( (args.batch_size, flow_loader_h, flow_loader_w)) # get grid for resampling XwarpP = 2 * ((XwarpP / (flow_loader_w - 1)) - 0.5) YwarpP = 2 * ((YwarpP / (flow_loader_h - 1)) - 0.5) gridP = torch.stack([XwarpP, YwarpP], dim=-1) # sampling with bilinear interpolation patch_var_past = torch.nn.functional.grid_sample(patch_var, gridP, mode="bilinear") mask_var_past = torch.nn.functional.grid_sample(mask_var, gridP, mode="bilinear") adv_tgt_img_var = torch.mul( (1 - mask_var), tgt_img_var) + torch.mul(mask_var, patch_var) adv_ref_past_img_var = torch.mul( (1 - mask_var_past), ref_past_img_var) + torch.mul( mask_var_past, patch_var_past) adv_ref_img_var = torch.mul( (1 - mask_var_future), ref_img_var) + torch.mul( mask_var_future, patch_var_future) adv_tgt_img_var = torch.clamp(adv_tgt_img_var, -1, 1) adv_ref_past_img_var = torch.clamp(adv_ref_past_img_var, -1, 1) adv_ref_img_var = torch.clamp(adv_ref_img_var, -1, 1) if type(flow_net).__name__ == 'Back2Future': adv_flow_fwd = flow_net(adv_ref_past_img_var, adv_tgt_img_var, adv_ref_img_var) else: adv_flow_fwd = flow_net(adv_tgt_img_var, adv_ref_img_var) # set patch to zero flow! mask_var_res = nn.functional.upsample(mask_var, size=(h_gt, w_gt), mode='bilinear') # Ignore patch motion if set! if args.ignore_mask_flow: forward_patch_flow = Variable(torch.cat((torch.zeros( (bt, 2, h_gt, w_gt)), torch.zeros((bt, 1, h_gt, w_gt))), 1).cuda(), volatile=True) flow_gt_var_adv = torch.mul( (1 - mask_var_res), flow_gt_var) + torch.mul( mask_var_res, forward_patch_flow) # import pdb; pdb.set_trace() epe = compute_epe(gt=flow_gt_var, pred=flow_fwd) adv_epe = compute_epe(gt=flow_gt_var_adv, pred=adv_flow_fwd) cos_sim = compute_cossim(flow_gt_var, flow_fwd) adv_cos_sim = compute_cossim(flow_gt_var_adv, adv_flow_fwd) errors.update([epe, adv_epe, cos_sim, adv_cos_sim]) if i % 1 == 0: index = i #int(i//10) imgs = normalize([tgt_img] + [ref_img_past] + [ref_img]) norm_tgt_img = imgs[0] norm_ref_img_past = imgs[1] norm_ref_img = imgs[2] patch_cpu = patch_var.data[0].cpu() mask_cpu = mask_var.data[0].cpu() adv_norm_tgt_img = normalize( adv_tgt_img_var.data.cpu() ) #torch.mul((1-mask_cpu), norm_tgt_img) + torch.mul(mask_cpu, patch_cpu) adv_norm_ref_img_past = normalize( adv_ref_past_img_var.data.cpu() ) # torch.mul((1-mask_cpu), norm_ref_img_past) + torch.mul(mask_cpu, patch_cpu) adv_norm_ref_img = normalize( adv_ref_img_var.data.cpu() ) #torch.mul((1-mask_cpu), norm_ref_img) + torch.mul(mask_cpu, patch_cpu) output_writer.add_image( 'val flow Input', transpose_image(tensor2array(norm_tgt_img[0])), 0) flow_to_show = flow_gt[0][:2, :, :].cpu() output_writer.add_image( 'val target Flow', transpose_image(flow_to_image(tensor2array(flow_to_show))), 0) # set flow to zero # zero_flow = Variable(torch.zeros(flow_fwd.shape).cuda(), volatile=True) # flow_fwd_masked = torch.mul((1-mask_var[:,:2,:,:]), flow_fwd) + torch.mul(mask_var[:,:2,:,:], zero_flow) flow_fwd_masked = flow_fwd # get ground truth flow val_GT_adv = flow_gt_var_adv.data[0].cpu().numpy().transpose( 1, 2, 0) # val_GT_adv = interp_gt_flow(val_GT_adv[:,:,:2], val_GT_adv[:,:,2]) val_GT_adv = cv2.resize(val_GT_adv, (flow_loader_w, flow_loader_h), interpolation=cv2.INTER_NEAREST) val_GT_adv[:, :, 0] = val_GT_adv[:, :, 0] * (flow_loader_w / w_gt) val_GT_adv[:, :, 1] = val_GT_adv[:, :, 1] * (flow_loader_h / h_gt) # gt normalization for visualization u = val_GT_adv[:, :, 0] v = val_GT_adv[:, :, 1] idxUnknow = (abs(u) > 1e7) | (abs(v) > 1e7) u[idxUnknow] = 0 v[idxUnknow] = 0 rad = np.sqrt(u**2 + v**2) maxrad = np.max(rad) val_GT_adv_Output = flow_to_image(val_GT_adv, maxrad) val_GT_adv_Output = cv2.erode(val_GT_adv_Output, np.ones((3, 3), np.uint8), iterations=1) # make points thicker val_GT_adv_Output = transpose_image(val_GT_adv_Output) / 255. val_Flow_Output = transpose_image( flow_to_image(tensor2array(flow_fwd.data[0].cpu()), maxrad)) / 255. val_adv_Flow_Output = transpose_image( flow_to_image(tensor2array(adv_flow_fwd.data[0].cpu()), maxrad)) / 255. val_Diff_Flow_Output = transpose_image( flow_to_image( tensor2array( (adv_flow_fwd - flow_fwd_masked).data[0].cpu()), maxrad)) / 255. val_tgt_image = transpose_image(tensor2array(norm_tgt_img[0])) val_ref_image = transpose_image(tensor2array(norm_ref_img[0])) val_adv_tgt_image = transpose_image( tensor2array(adv_norm_tgt_img[0])) val_adv_ref_image_past = transpose_image( tensor2array(adv_norm_ref_img_past[0])) val_adv_ref_image = transpose_image( tensor2array(adv_norm_ref_img[0])) val_patch = transpose_image(tensor2array(patch_var.data.cpu()[0])) # print(adv_norm_tgt_img.shape) # print(flow_fwd.data[0].cpu().shape) # if type(flow_net).__name__ == 'Back2Future': # val_output_viz = np.concatenate((val_adv_ref_image_past, val_adv_tgt_image, val_adv_ref_image, val_Flow_Output, val_adv_Flow_Output, val_Diff_Flow_Output), 2) # else: # val_output_viz = np.concatenate((val_adv_tgt_image, val_adv_ref_image, val_Flow_Output, val_adv_Flow_Output, val_Diff_Flow_Output, val_GT_adv_Output), 2) val_output_viz = np.concatenate( (val_ref_image, val_adv_ref_image, val_Flow_Output, val_adv_Flow_Output, val_Diff_Flow_Output, val_GT_adv_Output), 2) val_output_viz_im = Image.fromarray( (255 * val_output_viz.transpose(1, 2, 0)).astype('uint8')) val_output_viz_im.save(args.save_path / args.name + 'viz' + str(i).zfill(3) + '.jpg') output_writer.add_image('val Output viz {}'.format(index), val_output_viz, 0) #val_output_viz = np.vstack((val_Flow_Output, val_adv_Flow_Output, val_Diff_Flow_Output, val_adv_tgt_image, val_adv_ref_image)) #scipy.misc.imsave('outfile.jpg', os.path.join(output_vis_dir, 'vis_{}.png'.format(index))) result_scene_file.write( "{:10d}, {:10.4f}, {:10.4f}, {:10.4f}, {:10.4f}\n".format( i, epe, adv_epe, cos_sim, adv_cos_sim)) print("{:>10}, {:>10}, {:>10}, {:>10}".format(*error_names)) print("{:10.4f}, {:10.4f}, {:10.4f}, {:10.4f}".format(*errors.avg)) result_file.write( "{:10.4f}, {:10.4f}, {:10.4f}, {:10.4f}\n".format(*errors.avg)) result_scene_file.write( "{:>10}, {:10.4f}, {:10.4f}, {:10.4f}, {:10.4f}\n".format( *(["avg"] + errors.avg))) result_file.close() result_scene_file.close()
def main(): global args args = parser.parse_args() save_path = 'checkpoints/test_flownetc' if not os.path.exists(save_path): os.makedirs(save_path) summary_writer = SummaryWriter(save_path) normalize = custom_transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[1.0, 1.0, 1.0]) flow_loader_h, flow_loader_w = 384, 1280 valid_flow_transform = custom_transforms.Compose([ custom_transforms.Scale(h=flow_loader_h, w=flow_loader_w), custom_transforms.ArrayToTensor(), normalize ]) if args.dataset == "kitti2015": val_flow_set = ValidationFlowFlowNetC( root='/is/ps2/aranjan/AllFlowData/kitti/kitti2015', sequence_length=5, transform=valid_flow_transform) elif args.dataset == "kitti2012": val_flow_set = ValidationFlowKitti2012( root='/is/ps2/aranjan/AllFlowData/kitti/kitti2012', sequence_length=5, transform=valid_flow_transform) val_flow_loader = torch.utils.data.DataLoader(val_flow_set, batch_size=1, shuffle=False, num_workers=2, pin_memory=True, drop_last=True) flow_net = getattr(models, args.flownet)(pretrained=True).cuda() flow_net.eval() error_names = ['epe'] errors = AverageMeter(i=len(error_names)) for i, (tgt_img, ref_imgs, intrinsics, intrinsics_inv, flow_gt, flownet_c_flow, obj_map) in enumerate(val_flow_loader): tgt_img_var = Variable(tgt_img.cuda(), volatile=True) if args.dataset == "kitti2015": ref_imgs_var = [ Variable(img.cuda(), volatile=True) for img in ref_imgs ] ref_img_var = ref_imgs_var[2] elif args.dataset == "kitti2012": ref_img_var = Variable(ref_imgs.cuda(), volatile=True) flow_gt_var = Variable(flow_gt.cuda(), volatile=True) flownet_c_flow = Variable(flownet_c_flow.cuda(), volatile=True) # compute output flow_fwd = flow_net(tgt_img_var, ref_img_var) epe = compute_epe(gt=flownet_c_flow, pred=flow_fwd) scale_factor = compute_epe(gt=flownet_c_flow, pred=flow_fwd, op='div') #import ipdb #ipdb.set_trace() summary_writer.add_image('Frame 1', tensor2array(tgt_img_var.data[0].cpu()), i) summary_writer.add_image('Frame 2', tensor2array(ref_img_var.data[0].cpu()), i) summary_writer.add_image( 'Flow Output', flow_to_image(tensor2array(flow_fwd.data[0].cpu())), i) summary_writer.add_image( 'UnFlow Output', flow_to_image(tensor2array(flownet_c_flow.data[0][:2].cpu())), i) summary_writer.add_image( 'gtFlow Output', flow_to_image(tensor2array(flow_gt_var.data[0][:2].cpu())), i) summary_writer.add_image('EPE Image w UnFlow', tensor2array(epe.data.cpu()), i) summary_writer.add_scalar('EPE mean w UnFlow', epe.mean().data.cpu(), i) summary_writer.add_scalar('EPE max w UnFlow', epe.max().data.cpu(), i) summary_writer.add_scalar('Scale Factor max w UnFlow', scale_factor.max().data.cpu(), i) summary_writer.add_scalar('Scale Factor mean w UnFlow', scale_factor.mean().data.cpu(), i) summary_writer.add_scalar('Flow value max', flow_fwd.max().data.cpu(), i) print(i, "EPE: ", epe.mean().item()) #print(i, epe) #errors.update(epe) print('Done')