Load depth map (z in camera coordinate (x to right, y to down, z to far)) and surface normal in world coordinate (after bump map). Args: camera_name: frame_id: Returns: z_map: (h, w) z coordinate value in the camera coordinate
(
self,
camera_name: str,
frame_id: int,
)
| 425 | return hdf5_filename |
| 426 | |
| 427 | def _get_frame( |
| 428 | self, |
| 429 | camera_name: str, |
| 430 | frame_id: int, |
| 431 | ) -> T.Dict[str, torch.Tensor]: |
| 432 | """ |
| 433 | Load depth map (z in camera coordinate (x to right, y to down, z to far)) |
| 434 | and surface normal in world coordinate (after bump map). |
| 435 | |
| 436 | Args: |
| 437 | camera_name: |
| 438 | frame_id: |
| 439 | |
| 440 | Returns: |
| 441 | z_map: (h, w) z coordinate value in the camera coordinate |
| 442 | surface_normal_w: (h, w, 3) surface normal in world coordinate |
| 443 | rgb: (h, w, 3) color before tone mapping |
| 444 | """ |
| 445 | |
| 446 | # calculate z_map |
| 447 | filename = self._get_geometry_hdf5_filename( |
| 448 | type='depth_meters', |
| 449 | camera_name=camera_name, |
| 450 | frame_id=frame_id, |
| 451 | ) |
| 452 | with h5py.File(filename, "r") as f: |
| 453 | ray_ts = f["dataset"][:] # (h, w) in meter |
| 454 | ray_ts_abs = torch.from_numpy(ray_ts.astype(np.float32)) # (h, w) in meter |
| 455 | |
| 456 | u_min, u_max, v_min, v_max = -1.0, 1.0, -1.0, 1.0 |
| 457 | half_du = 0.5 * (u_max - u_min) / self.width_px |
| 458 | half_dv = 0.5 * (v_max - v_min) / self.height_px |
| 459 | u, v = torch.meshgrid( |
| 460 | torch.linspace(u_min + half_du, u_max - half_du, self.width_px), |
| 461 | torch.linspace(v_min + half_dv, v_max - half_dv, self.height_px).flip(dims=[0]), |
| 462 | indexing='xy', |
| 463 | ) # (h, w) # (h, w) |
| 464 | uv1 = torch.stack((u, v, torch.ones_like(u)), dim=-1) # (h, w, 3) |
| 465 | uv1 = uv1.reshape(-1, 3) # (hw, 3) |
| 466 | |
| 467 | # compute our own rays |
| 468 | H_c2w = self.H_c2ws[camera_name][self.cam_name_to_frame_idxs_to_idxs[camera_name][frame_id]] # (4, 4) |
| 469 | tmp_c = (self.M_u2c @ uv1.T).T # (hw, 3) |
| 470 | tmp_w = (H_c2w[:3, :3] @ tmp_c.T).T # (hw, 3) |
| 471 | ray_direction_w = torch.nn.functional.normalize(tmp_w, p=2, dim=-1) # (hw, 3) |
| 472 | xyz_w = H_c2w[:3, 3].reshape(1, 3) + ray_direction_w * ray_ts_abs.reshape(-1, 1) # (hw, 3) |
| 473 | |
| 474 | H_w2c = rigid_motion.inv_homogeneous_tensors(H_c2w) # (4, 4) |
| 475 | xyz1_w = torch.cat((xyz_w, torch.ones_like(xyz_w[:, :1])), dim=-1) # (hw, 4) |
| 476 | xyz1_c = (H_w2c @ xyz1_w.T).T # (hw, 4) |
| 477 | xyz_c = xyz1_c[:, :3] # (hw, 3) |
| 478 | xyz_c = xyz_c.reshape(self.height_px, self.width_px, 3) # (h, w, 3) |
| 479 | z_map = xyz_c[..., 2] # (h, w) all < 0 |
| 480 | z_map = -1 * z_map # (h, w) all >= 0 |
| 481 | |
| 482 | # surface_normal_w |
| 483 | filename = self._get_geometry_hdf5_filename( |
| 484 | type='normal_bump_world', |
no test coverage detected