Compute the zdir and dps in the world coordinate using z_map and camera pose. Args: z_map: (*, h, w) the z coordinate of the point in the camera coordinate on the sensor, not along the corresponding camera ray. intrinsic: (*, 3, 3) camera
(
z_map: torch.Tensor,
intrinsic: torch.Tensor,
H_c2w: torch.Tensor,
subsample: int = 1,
)
| 1370 | |
| 1371 | |
| 1372 | def compute_3d_zdir_and_dps( |
| 1373 | z_map: torch.Tensor, |
| 1374 | intrinsic: torch.Tensor, |
| 1375 | H_c2w: torch.Tensor, |
| 1376 | subsample: int = 1, |
| 1377 | ): |
| 1378 | """ |
| 1379 | Compute the zdir and dps in the world coordinate using z_map and camera pose. |
| 1380 | |
| 1381 | Args: |
| 1382 | z_map: |
| 1383 | (*, h, w) the z coordinate of the point in the camera coordinate on the sensor, |
| 1384 | not along the corresponding camera ray. |
| 1385 | intrinsic: |
| 1386 | (*, 3, 3) camera intrinsic matrix |
| 1387 | H_c2w: |
| 1388 | (*, 4, 4) homegeneous matrix that convert camera coord to world coord. |
| 1389 | Note that the y axis should be inverted in the cam_poses. |
| 1390 | subsample: int |
| 1391 | index stride |
| 1392 | other_maps: |
| 1393 | a list of (*, h, w, d) to associated with each point |
| 1394 | |
| 1395 | Returns: |
| 1396 | zdir_w: |
| 1397 | (*, h//subsample, w//subsample, 3) camera z direction in world coordinates |
| 1398 | dps_w: |
| 1399 | (*, h//subsample, w//subsample, 1) distance per sample in world coordinates |
| 1400 | dps_uw: |
| 1401 | (*, h//subsample, w//subsample, 3) distance per sample in u direction in world coordinates |
| 1402 | dps_vw: |
| 1403 | (*, h//subsample, w//subsample, 3) distance per sample in v direction in world coordinates |
| 1404 | """ |
| 1405 | |
| 1406 | dtype = z_map.dtype |
| 1407 | device = z_map.device |
| 1408 | h, w = z_map.size(-2), z_map.size(-1) |
| 1409 | |
| 1410 | # generate u v w on the sensor coord |
| 1411 | u, v = torch.meshgrid( |
| 1412 | torch.arange(0, w, subsample, device=device), |
| 1413 | torch.arange(0, h, subsample, device=device), |
| 1414 | indexing='xy', |
| 1415 | ) # u: (h', w') for x, v: (h', w') for y in the sensor coord |
| 1416 | |
| 1417 | z_map = z_map[..., v, u] # (*, h', w') |
| 1418 | valid_z = z_map < 1e11 # default background z is set to 1e12 |
| 1419 | inv_intrinsic = torch.linalg.inv(intrinsic).unsqueeze(-3).unsqueeze(-3) # (*, 1, 1, 3, 3) |
| 1420 | |
| 1421 | # get distance per sample in u/v direction; |
| 1422 | dps_u = torch.stack((subsample * z_map, 0 * z_map), dim=-1) # (*, h, w, 2) # distance per sample in u direction |
| 1423 | dps_v = torch.stack((0 * z_map, subsample * z_map), dim=-1) # (*, h, w, 2) # distance per sample in v direction |
| 1424 | dps_uv = torch.stack((dps_u, dps_v), dim=-1) |
| 1425 | |
| 1426 | dps_uvc = inv_intrinsic[..., :2, :2] @ dps_uv # (*, h', w', 2, 2) distance in cam cord |
| 1427 | |
| 1428 | dps_uvc_shape = list(dps_uvc.shape) |
| 1429 | dps_uvc_shape[-2] = 1 |