File size: 4,016 Bytes
92011e4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 | # Adapted from TencentARC/Pixal3D, Copyright (c) 2026 Tencent, MIT (see LICENSE-PIXAL3D).
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)
|