Args: rgbs: [S C H W] depths_gt: [S C H W] poses_gt: [S C] poses_pred: [S C] depth_pred: [S H W]
(rgbs, poses_pred, poses_gt,
depth_gt, depth_pred, iter_num=0,
vis=None, logger_tf=None, cfg=None)
| 848 | return pcl |
| 849 | |
| 850 | def vis_result(rgbs, poses_pred, poses_gt, |
| 851 | depth_gt, depth_pred, iter_num=0, |
| 852 | vis=None, logger_tf=None, cfg=None): |
| 853 | """ |
| 854 | Args: |
| 855 | rgbs: [S C H W] |
| 856 | depths_gt: [S C H W] |
| 857 | poses_gt: [S C] |
| 858 | poses_pred: [S C] |
| 859 | depth_pred: [S H W] |
| 860 | """ |
| 861 | assert len(rgbs.shape) == 4, "only support one sequence, T 3 H W of rbg" |
| 862 | |
| 863 | if vis is None: |
| 864 | return |
| 865 | S, _, H, W = depth_gt.shape |
| 866 | # get the xy |
| 867 | yx = torch.meshgrid(torch.arange(H).to(depth_pred.device), |
| 868 | torch.arange(W).to(depth_pred.device),indexing='ij') |
| 869 | xy = torch.stack(yx[::-1], dim=0).float().to(depth_pred.device) |
| 870 | xy_norm = (xy / torch.tensor([W, H], |
| 871 | device=depth_pred.device).view(2, 1, 1) - 0.5)*2 |
| 872 | xy = xy[None].repeat(S, 1, 1, 1) |
| 873 | xy_depth = torch.cat([xy, depth_pred[:,None]], dim=1).permute(0, 2, 3, 1) |
| 874 | xy_depth_gt = torch.cat([xy, depth_gt], dim=1).permute(0, 2, 3, 1) |
| 875 | # get the focal length |
| 876 | focal_length = poses_gt[:,-1]*max(H, W) |
| 877 | |
| 878 | # vis the camera poses |
| 879 | poses_gt_vis = pose_encoding_to_camera(poses_gt, |
| 880 | pose_encoding_type="absT_quaR_OneFL",to_OpenCV=False) |
| 881 | poses_pred_vis = pose_encoding_to_camera(poses_pred, |
| 882 | pose_encoding_type="absT_quaR_OneFL",to_OpenCV=False) |
| 883 | |
| 884 | R_gt = poses_gt_vis.R.float() |
| 885 | R_pred = poses_pred_vis.R.float() |
| 886 | T_gt = poses_gt_vis.T.float() |
| 887 | T_pred = poses_pred_vis.T.float() |
| 888 | # C2W poses |
| 889 | R_gt_c2w = R_gt.permute(0,2,1) |
| 890 | T_gt_c2w = (-R_gt_c2w @ T_gt[:, :, None]).squeeze(-1) |
| 891 | R_pred_c2w = R_pred.permute(0,2,1) |
| 892 | T_pred_c2w = (-R_pred_c2w @ T_pred[:, :, None]).squeeze(-1) |
| 893 | with torch.cuda.amp.autocast(enabled=False): |
| 894 | pick_idx = torch.randperm(S)[:min(24, S)] |
| 895 | # pick_idx = [1] |
| 896 | #NOTE: very strange that the camera need C2W Rotation and W2C translation as input |
| 897 | poses_gt_vis = PerspectiveCamerasVisual( |
| 898 | R=R_gt_c2w[pick_idx], T=T_gt[pick_idx], |
| 899 | device=poses_gt_vis.device, image_size=((H, W),) |
| 900 | ) |
| 901 | poses_pred_vis = PerspectiveCamerasVisual( |
| 902 | R=R_pred_c2w[pick_idx], T=T_pred[pick_idx], |
| 903 | device=poses_pred_vis.device |
| 904 | ) |
| 905 | visual_dict = {"scenes": {"cameras": poses_pred_vis, "cameras_gt": poses_gt_vis}} |
| 906 | env_name = f"train_visualize_iter_{iter_num:05d}" |
| 907 | print(f"Visualizing the scene by visdom at env: {env_name}") |
nothing calls this directly
no test coverage detected