| |
|
|
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
|
|
|
|
| def project_points_to_image_batch( |
| points_3d: torch.Tensor, |
| transform_matrix: torch.Tensor, |
| camera_angle_x: torch.Tensor, |
| resolution: int, |
| ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| """Project 3D points to pixel coordinates with a pinhole camera (32 mm sensor width). |
| |
| Args: |
| points_3d: B N 3 points. |
| transform_matrix: B 4 4 camera-to-world matrices. |
| camera_angle_x: B horizontal fields of view, radians. |
| resolution: Square image resolution in pixels. |
| |
| Returns: |
| Pixel coordinates (B N 2), depth (B N) and a mask of points inside the view (B N). |
| """ |
| batch, count = points_3d.shape[:2] |
| ones = torch.ones(batch, count, 1, device=points_3d.device, dtype=points_3d.dtype) |
| points_homogeneous = torch.cat([points_3d, ones], dim=-1) |
| world_to_camera = torch.linalg.inv(transform_matrix.float()).to(transform_matrix.dtype) |
| points_camera = torch.bmm(points_homogeneous, world_to_camera.transpose(-2, -1))[..., :3] |
| x_cam, y_cam, z_cam = points_camera.unbind(-1) |
| depth = -z_cam |
| sensor_width = 32.0 |
| focal_length = 16.0 / torch.tan(camera_angle_x / 2.0) |
| focal_length_pixels = (focal_length * resolution / sensor_width).unsqueeze(1) |
| x_ndc = focal_length_pixels * x_cam / (-z_cam + 1e-08) |
| y_ndc = focal_length_pixels * y_cam / (-z_cam + 1e-08) |
| x_pixel = x_ndc + resolution / 2.0 |
| y_pixel = -y_ndc + resolution / 2.0 |
| valid_mask = (x_pixel >= 0) & (x_pixel < resolution) & (y_pixel >= 0) & (y_pixel < resolution) & (depth > 0) |
| return torch.stack([x_pixel, y_pixel], dim=-1), depth, valid_mask |
|
|
|
|
| class ProjGrid(nn.Module): |
| """Project a 3D grid (or selected grid cells) into the front view and sample image features there.""" |
|
|
| def __init__(self, grid_resolution: int, image_resolution: int): |
| super().__init__() |
| self.image_resolution = image_resolution |
| one_dim = torch.linspace(-1, 1, grid_resolution) |
| x, y, z = torch.meshgrid(one_dim, one_dim, one_dim, indexing="ij") |
| rotation_matrix = torch.tensor([[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]]) |
| grid_points = torch.matmul(torch.stack((x, y, z), dim=-1), rotation_matrix.T).reshape(-1, 3) |
| self.register_buffer("grid_points", grid_points) |
| front_view_transform_matrix = torch.tensor( |
| [ |
| [1.0, 0.0, 0.0, 0.0], |
| [0.0, 0.0, -1.0, -2.0], |
| [0.0, 1.0, 0.0, 0.0], |
| [0.0, 0.0, 0.0, 1.0], |
| ] |
| ) |
| self.register_buffer("front_view_transform_matrix", front_view_transform_matrix) |
|
|
| def forward( |
| self, |
| features_map: torch.Tensor, |
| camera_angle_x: torch.Tensor, |
| distance: torch.Tensor, |
| mesh_scale: torch.Tensor, |
| grid_indices: torch.Tensor | None = None, |
| ) -> torch.Tensor: |
| """Sample a B C H W feature map at the projected grid points; returns B N C.""" |
| batch = features_map.shape[0] |
| grid_points = self.grid_points |
| if grid_indices is not None: |
| grid_points = grid_points[grid_indices] |
| grid_points = grid_points.expand(batch, -1, -1) |
| grid_points = grid_points / mesh_scale.unsqueeze(-1).unsqueeze(-1) / 2 |
| transform_matrix = self.front_view_transform_matrix.expand(batch, -1, -1).clone() |
| transform_matrix[:, 1, 3] = -distance |
| image_points, _, _ = project_points_to_image_batch( |
| grid_points, transform_matrix, camera_angle_x, self.image_resolution |
| ) |
| image_points_norm = (image_points + 0.5) / self.image_resolution * 2 - 1 |
| sample_grid = image_points_norm.view(batch, -1, 1, 2) |
| sampled = F.grid_sample(features_map, sample_grid, mode="bilinear", align_corners=False, padding_mode="border") |
| return sampled.squeeze(-1).permute(0, 2, 1) |
|
|
