|
"""Contain the functions to sample point features from the triplane |
|
representation.""" |
|
|
|
import torch |
|
|
|
__all__ = ['TriplaneSampler'] |
|
|
|
|
|
class TriplaneSampler(torch.nn.Module): |
|
"""Defines the class to help sample point features from the triplane |
|
representation. |
|
|
|
Basically, this class implements the following functions for sampling point |
|
features (rgb && sigma) from the triplane representation: |
|
|
|
1. `generate_planes()`. |
|
2. `project_onto_planes()`. |
|
3. `sample_from_planes()`. |
|
4. `sample_from_3dgrid()`. |
|
""" |
|
|
|
def __init__(self): |
|
super().__init__() |
|
|
|
@staticmethod |
|
def generate_planes(): |
|
""" |
|
Defines planes by the three vectors that form the "axes" of the |
|
plane. Should work with arbitrary number of planes and planes of |
|
arbitrary orientation. |
|
""" |
|
return torch.tensor([[[1, 0, 0], |
|
[0, 1, 0], |
|
[0, 0, 1]], |
|
[[1, 0, 0], |
|
[0, 0, 1], |
|
[0, 1, 0]], |
|
[[0, 0, 1], |
|
[1, 0, 0], |
|
[0, 1, 0]]], dtype=torch.float32) |
|
|
|
@staticmethod |
|
def project_onto_planes(planes, coordinates): |
|
""" |
|
Does a projection of a 3D point onto a batch of 2D planes, |
|
returning 2D plane coordinates. |
|
|
|
Args: |
|
planes: Plane axes of shape (n_planes, 3, 3) |
|
coordinates: Coordinates of shape (N, M, 3) |
|
|
|
Returns: |
|
projections: Projections of shape (N*n_planes, M, 2) |
|
""" |
|
N, M, C = coordinates.shape |
|
n_planes, _, _ = planes.shape |
|
coordinates = coordinates.unsqueeze(1).expand(-1, n_planes, -1, |
|
-1).reshape( |
|
N * n_planes, M, 3) |
|
inv_planes = torch.linalg.inv(planes).unsqueeze(0).expand( |
|
N, -1, -1, -1).reshape(N * n_planes, 3, 3) |
|
projections = torch.bmm(coordinates, inv_planes) |
|
return projections[..., :2] |
|
|
|
@staticmethod |
|
def sample_from_planes(plane_axes, |
|
plane_features, |
|
coordinates, |
|
mode='bilinear', |
|
padding_mode='zeros', |
|
box_warp=None): |
|
assert padding_mode == 'zeros' |
|
N, n_planes, C, H, W = plane_features.shape |
|
_, M, _ = coordinates.shape |
|
plane_features = plane_features.view(N * n_planes, C, H, W) |
|
|
|
coordinates = (2 / box_warp) * coordinates |
|
|
|
projected_coordinates = TriplaneSampler.project_onto_planes( |
|
plane_axes, coordinates).unsqueeze(1) |
|
output_features = torch.nn.functional.grid_sample( |
|
plane_features, |
|
projected_coordinates.float(), |
|
mode=mode, |
|
padding_mode=padding_mode, |
|
align_corners=False).permute(0, 3, 2, |
|
1).reshape(N, n_planes, M, C) |
|
return output_features |
|
|
|
@staticmethod |
|
def sample_from_3dgrid(grid, coordinates): |
|
""" |
|
Expects coordinates in shape (batch_size, num_points_per_batch, 3) |
|
Expects grid in shape (1, channels, H, W, D) |
|
(Also works if grid has batch size) |
|
Returns: |
|
Sampled features |
|
with shape: (batch_size, num_points_per_batch, feature_channels). |
|
""" |
|
batch_size, n_coords, n_dims = coordinates.shape |
|
sampled_features = torch.nn.functional.grid_sample( |
|
grid.expand(batch_size, -1, -1, -1, -1), |
|
coordinates.reshape(batch_size, 1, 1, -1, n_dims), |
|
mode='bilinear', |
|
padding_mode='zeros', |
|
align_corners=False) |
|
N, C, H, W, D = sampled_features.shape |
|
sampled_features = sampled_features.permute(0, 4, 3, 2, 1).reshape( |
|
N, H * W * D, C) |
|
return sampled_features |
