args: xy_depth: [S H W 3] focal_length: [S] R: [S 3 3] W2C T: [S 3] W2C return: xyz: [S 3 (H W)]
(
xy_depth: torch.Tensor,
focal_length: torch.Tensor,
R: torch.Tensor,
T: torch.Tensor,
device: torch.device = None,
H: int = 518,
W: int = 518
)
| 926 | return |
| 927 | |
| 928 | def depth2pcd( |
| 929 | xy_depth: torch.Tensor, |
| 930 | focal_length: torch.Tensor, |
| 931 | R: torch.Tensor, |
| 932 | T: torch.Tensor, |
| 933 | device: torch.device = None, |
| 934 | H: int = 518, |
| 935 | W: int = 518 |
| 936 | ): |
| 937 | """ |
| 938 | args: |
| 939 | xy_depth: [S H W 3] |
| 940 | focal_length: [S] |
| 941 | R: [S 3 3] W2C |
| 942 | T: [S 3] W2C |
| 943 | return: |
| 944 | xyz: [S 3 (H W)] |
| 945 | """ |
| 946 | S, H, W, _ = xy_depth.shape |
| 947 | # get the intrinsic |
| 948 | K = torch.eye(3, device=device)[None].repeat(len(focal_length), 1, 1).to(device) |
| 949 | K[:, 0, 0] = focal_length |
| 950 | K[:, 1, 1] = focal_length |
| 951 | K[:, 0, 2] = 0.5 * W |
| 952 | K[:, 1, 2] = 0.5 * H |
| 953 | K_inv = K.inverse() |
| 954 | # xyz |
| 955 | xyz = xy_depth.view(S, -1, 3).permute(0, 2, 1) # S 3 (H W) |
| 956 | depth = xyz[:, 2:].clone() # S (H W) 1 |
| 957 | xyz[:, 2] = 1 |
| 958 | xyz = K_inv @ xyz # S 3 (H W) |
| 959 | xyz = xyz * depth |
| 960 | # to world coordinate |
| 961 | xyz = R.permute(0,2,1) @ (xyz - T[:, :, None]) |
| 962 | |
| 963 | return xyz |
| 964 | |
| 965 | |
| 966 | def pose_enc2mat(poses_pred, |