def _get_ndc_grid(self, h, w, device):
if w >= h:
range_x = w / h
range_y = 1.0
else:
range_x = 1.0
range_y = h / w
half_pix_width = range_x / w
half_pix_height = range_y / h
min_x = range_x - half_pix_width
max_x = -range_x + half_pix_width
min_y = range_y - half_pix_height
max_y = -range_y + half_pix_height
y_grid, x_grid = meshgrid_ij(
torch.linspace(min_y, max_y, h, dtype=torch.float32),
torch.linspace(min_x, max_x, w, dtype=torch.float32),
)
x_points = x_grid.contiguous().view(-1).to(device)
y_points = y_grid.contiguous().view(-1).to(device)
xy = torch.stack((x_points, y_points), dim=1)
return xy
def __init__(
self,
*,
min_x: float,
max_x: float,
min_y: float,
max_y: float,
image_width: int,
image_height: int,
n_pts_per_ray: int,
min_depth: float,
max_depth: float,
n_rays_per_image: Optional[int] = None,
unit_directions: bool = False,
stratified_sampling: bool = False,
) -> None:
"""
Args:
min_x: The leftmost x-coordinate of each ray's source pixel's center.
max_x: The rightmost x-coordinate of each ray's source pixel's center.
min_y: The topmost y-coordinate of each ray's source pixel's center.
max_y: The bottommost y-coordinate of each ray's source pixel's center.
image_width: The horizontal size of the image grid.
image_height: The vertical size of the image grid.
n_pts_per_ray: The number of points sampled along each ray.
min_depth: The minimum depth of a ray-point.
max_depth: The maximum depth of a ray-point.
n_rays_per_image: If given, this amount of rays are sampled from the grid.
unit_directions: whether to normalize direction vectors in ray bundle.
stratified_sampling: if set, performs stratified random sampling
along the ray; otherwise takes ray points at deterministic offsets.
"""
super().__init__()
self._n_pts_per_ray = n_pts_per_ray
self._min_depth = min_depth
self._max_depth = max_depth
self._n_rays_per_image = n_rays_per_image
self._unit_directions = unit_directions
self._stratified_sampling = stratified_sampling
# get the initial grid of image xy coords
_xy_grid = torch.stack(
tuple(
reversed(
meshgrid_ij(
torch.linspace(min_y, max_y, image_height, dtype=torch.float32),
torch.linspace(min_x, max_x, image_width, dtype=torch.float32),
)
)
),
dim=-1,
)
self.register_buffer("_xy_grid", _xy_grid, persistent=False)
def _calculate_coordinate_grid(self,
world_coordinates: bool = True
) -> torch.Tensor:
"""
Calculate the 3D coordinate grid of the volumetric grid either in
in local (`world_coordinates=False`) or
world coordinates (`world_coordinates=True`) .
"""
densities = self.densities()
ba, _, de, he, wi = densities.shape
grid_sizes = self.get_grid_sizes()
# generate coordinate axes
vol_axes = [
torch.linspace(-1.0,
1.0,
r,
dtype=torch.float32,
device=self.device) for r in (de, he, wi)
]
# generate per-coord meshgrids
Z, Y, X = meshgrid_ij(vol_axes)
# stack the coord grids ... this order matches the coordinate convention
# of torch.nn.grid_sample
vol_coords_local = torch.stack((X, Y, Z),
dim=3)[None].repeat(ba, 1, 1, 1, 1)
# get grid sizes relative to the maximal volume size
grid_sizes_relative = (torch.tensor(
[[de, he, wi]], device=grid_sizes.device, dtype=torch.float32) -
1) / (grid_sizes - 1).float()
if (grid_sizes_relative != 1.0).any():
# if any of the relative sizes != 1.0, adjust the grid
grid_sizes_relative_reshape = grid_sizes_relative[:,
[2, 1, 0]][:,
None,
None,
None]
vol_coords_local *= grid_sizes_relative_reshape
vol_coords_local += grid_sizes_relative_reshape - 1
if world_coordinates:
vol_coords = self.local_to_world_coords(vol_coords_local)
else:
vol_coords = vol_coords_local
return vol_coords
def cubify(voxels, thresh, device=None, align: str = "topleft") -> Meshes:
r"""
Converts a voxel to a mesh by replacing each occupied voxel with a cube
consisting of 12 faces and 8 vertices. Shared vertices are merged, and
internal faces are removed.
Args:
voxels: A FloatTensor of shape (N, D, H, W) containing occupancy probabilities.
thresh: A scalar threshold. If a voxel occupancy is larger than
thresh, the voxel is considered occupied.
device: The device of the output meshes
align: Defines the alignment of the mesh vertices and the grid locations.
Has to be one of {"topleft", "corner", "center"}. See below for explanation.
Default is "topleft".
Returns:
meshes: A Meshes object of the corresponding meshes.
The alignment between the vertices of the cubified mesh and the voxel locations (or pixels)
is defined by the choice of `align`. We support three modes, as shown below for a 2x2 grid:
X---X---- X-------X ---------
| | | | | | | X | X |
X---X---- --------- ---------
| | | | | | | X | X |
--------- X-------X ---------
topleft corner center
In the figure, X denote the grid locations and the squares represent the added cuboids.
When `align="topleft"`, then the top left corner of each cuboid corresponds to the
pixel coordinate of the input grid.
When `align="corner"`, then the corners of the output mesh span the whole grid.
When `align="center"`, then the grid locations form the center of the cuboids.
"""
if device is None:
device = voxels.device
if align not in ["topleft", "corner", "center"]:
raise ValueError(
"Align mode must be one of (topleft, corner, center).")
if len(voxels) == 0:
return Meshes(verts=[], faces=[])
N, D, H, W = voxels.size()
# vertices corresponding to a unit cube: 8x3
cube_verts = torch.tensor(
[
[0, 0, 0],
[0, 0, 1],
[0, 1, 0],
[0, 1, 1],
[1, 0, 0],
[1, 0, 1],
[1, 1, 0],
[1, 1, 1],
],
dtype=torch.int64,
device=device,
)
# faces corresponding to a unit cube: 12x3
cube_faces = torch.tensor(
[
[0, 1, 2],
[1, 3, 2], # left face: 0, 1
[2, 3, 6],
[3, 7, 6], # bottom face: 2, 3
[0, 2, 6],
[0, 6, 4], # front face: 4, 5
[0, 5, 1],
[0, 4, 5], # up face: 6, 7
[6, 7, 5],
[6, 5, 4], # right face: 8, 9
[1, 7, 3],
[1, 5, 7], # back face: 10, 11
],
dtype=torch.int64,
device=device,
)
wx = torch.tensor([0.5, 0.5], device=device).view(1, 1, 1, 1, 2)
wy = torch.tensor([0.5, 0.5], device=device).view(1, 1, 1, 2, 1)
wz = torch.tensor([0.5, 0.5], device=device).view(1, 1, 2, 1, 1)
voxelt = voxels.ge(thresh).float()
# N x 1 x D x H x W
voxelt = voxelt.view(N, 1, D, H, W)
# N x 1 x (D-1) x (H-1) x (W-1)
voxelt_x = F.conv3d(voxelt, wx).gt(0.5).float()
voxelt_y = F.conv3d(voxelt, wy).gt(0.5).float()
voxelt_z = F.conv3d(voxelt, wz).gt(0.5).float()
# 12 x N x 1 x D x H x W
faces_idx = torch.ones((cube_faces.size(0), N, 1, D, H, W), device=device)
# add left face
faces_idx[0, :, :, :, :, 1:] = 1 - voxelt_x
faces_idx[1, :, :, :, :, 1:] = 1 - voxelt_x
# add bottom face
faces_idx[2, :, :, :, :-1, :] = 1 - voxelt_y
faces_idx[3, :, :, :, :-1, :] = 1 - voxelt_y
# add front face
faces_idx[4, :, :, 1:, :, :] = 1 - voxelt_z
faces_idx[5, :, :, 1:, :, :] = 1 - voxelt_z
# add up face
faces_idx[6, :, :, :, 1:, :] = 1 - voxelt_y
faces_idx[7, :, :, :, 1:, :] = 1 - voxelt_y
# add right face
faces_idx[8, :, :, :, :, :-1] = 1 - voxelt_x
faces_idx[9, :, :, :, :, :-1] = 1 - voxelt_x
# add back face
faces_idx[10, :, :, :-1, :, :] = 1 - voxelt_z
faces_idx[11, :, :, :-1, :, :] = 1 - voxelt_z
faces_idx *= voxelt
# N x H x W x D x 12
faces_idx = faces_idx.permute(1, 2, 4, 5, 3, 0).squeeze(1)
# (NHWD) x 12
faces_idx = faces_idx.contiguous()
faces_idx = faces_idx.view(-1, cube_faces.size(0))
# boolean to linear index
# NF x 2
linind = torch.nonzero(faces_idx, as_tuple=False)
# NF x 4
nyxz = unravel_index(linind[:, 0], (N, H, W, D))
# NF x 3: faces
faces = torch.index_select(cube_faces, 0, linind[:, 1])
grid_faces = []
for d in range(cube_faces.size(1)):
# NF x 3
xyz = torch.index_select(cube_verts, 0, faces[:, d])
permute_idx = torch.tensor([1, 0, 2], device=device)
yxz = torch.index_select(xyz, 1, permute_idx)
yxz += nyxz[:, 1:]
# NF x 1
temp = ravel_index(yxz, (H + 1, W + 1, D + 1))
grid_faces.append(temp)
# NF x 3
grid_faces = torch.stack(grid_faces, dim=1)
y, x, z = meshgrid_ij(torch.arange(H + 1), torch.arange(W + 1),
torch.arange(D + 1))
y = y.to(device=device, dtype=torch.float32)
x = x.to(device=device, dtype=torch.float32)
z = z.to(device=device, dtype=torch.float32)
if align == "center":
x = x - 0.5
y = y - 0.5
z = z - 0.5
margin = 0.0 if align == "corner" else 1.0
y = y * 2.0 / (H - margin) - 1.0
x = x * 2.0 / (W - margin) - 1.0
z = z * 2.0 / (D - margin) - 1.0
# ((H+1)(W+1)(D+1)) x 3
grid_verts = torch.stack((x, y, z), dim=3).view(-1, 3)
if len(nyxz) == 0:
verts_list = [torch.tensor([], dtype=torch.float32, device=device)] * N
faces_list = [torch.tensor([], dtype=torch.int64, device=device)] * N
return Meshes(verts=verts_list, faces=faces_list)
num_verts = grid_verts.size(0)
grid_faces += nyxz[:, 0].view(-1, 1) * num_verts
idleverts = torch.ones(num_verts * N, dtype=torch.uint8, device=device)
indices = grid_faces.flatten()
if device.type == "cpu":
indices = torch.unique(indices)
idleverts.scatter_(0, indices, 0)
grid_faces -= nyxz[:, 0].view(-1, 1) * num_verts
split_size = torch.bincount(nyxz[:, 0], minlength=N)
faces_list = list(torch.split(grid_faces, split_size.tolist(), 0))
idleverts = idleverts.view(N, num_verts)
idlenum = idleverts.cumsum(1)
verts_list = [
# pyre-fixme[16]: `Tensor` has no attribute `index_select`.
grid_verts.index_select(0,
(idleverts[n] == 0).nonzero(as_tuple=False)[:,
0])
for n in range(N)
]
faces_list = [
nface - idlenum[n][nface] for n, nface in enumerate(faces_list)
]
return Meshes(verts=verts_list, faces=faces_list)
def test_ndc_grid_sample_rendering(self):
"""
Use PyTorch3D point renderer to render a colored point cloud, then
sample the image at the locations of the point projections with
`ndc_grid_sample`. Finally, assert that the sampled colors are equal to the
original point cloud colors.
Note that, in order to ensure correctness, we use a nearest-neighbor
assignment point renderer (i.e. no soft splatting).
"""
# generate a bunch of 3D points on a regular grid lying in the z-plane
n_grid_pts = 10
grid_scale = 0.9
z_plane = 2.0
image_size = [128, 128]
point_radius = 0.015
n_pts = n_grid_pts * n_grid_pts
pts = torch.stack(
meshgrid_ij([torch.linspace(-grid_scale, grid_scale, n_grid_pts)] *
2, ),
dim=-1,
)
pts = torch.cat([pts, z_plane * torch.ones_like(pts[..., :1])], dim=-1)
pts = pts.reshape(1, n_pts, 3)
# color the points randomly
pts_colors = torch.rand(1, n_pts, 3)
# make trivial rendering cameras
cameras = PerspectiveCameras(
R=eyes(dim=3, N=1),
device=pts.device,
T=torch.zeros(1, 3, dtype=torch.float32, device=pts.device),
)
# render the point cloud
pcl = Pointclouds(points=pts, features=pts_colors)
renderer = NearestNeighborPointsRenderer(
rasterizer=PointsRasterizer(
cameras=cameras,
raster_settings=PointsRasterizationSettings(
image_size=image_size,
radius=point_radius,
points_per_pixel=1,
),
),
compositor=AlphaCompositor(),
)
im_render = renderer(pcl)
# sample the render at projected pts
pts_proj = cameras.transform_points(pcl.points_padded())[..., :2]
pts_colors_sampled = ndc_grid_sample(
im_render,
pts_proj,
mode="nearest",
align_corners=False,
).permute(0, 2, 1)
# assert that the samples are the same as original points
self.assertClose(pts_colors, pts_colors_sampled, atol=1e-4)
def make_material_atlas(image: torch.Tensor, faces_verts_uvs: torch.Tensor,
texture_size: int) -> torch.Tensor:
r"""
Given a single texture image and the uv coordinates for all the
face vertices, create a square texture map per face using
the formulation from [1].
For a triangle with vertices (v0, v1, v2) we can create a barycentric coordinate system
with the x axis being the vector (v0 - v2) and the y axis being the vector (v1 - v2).
The barycentric coordinates range from [0, 1] in the +x and +y direction so this creates
a triangular texture space with vertices at (0, 1), (0, 0) and (1, 0).
The per face texture map is of shape (texture_size, texture_size, 3)
which is a square. To map a triangular texture to a square grid, each
triangle is parametrized as follows (e.g. R = texture_size = 3):
The triangle texture is first divided into RxR = 9 subtriangles which each
map to one grid cell. The numbers in the grid cells and triangles show the mapping.
..code-block::python
Triangular Texture Space:
1
|\
|6 \
|____\
|\ 7 |\
|3 \ |4 \
|____\|____\
|\ 8 |\ 5 |\
|0 \ |1 \ |2 \
|____\|____\|____\
0 1
Square per face texture map:
R ____________________
| | | |
| 6 | 7 | 8 |
|______|______|______|
| | | |
| 3 | 4 | 5 |
|______|______|______|
| | | |
| 0 | 1 | 2 |
|______|______|______|
0 R
The barycentric coordinates of each grid cell are calculated using the
xy coordinates:
..code-block::python
The cartesian coordinates are:
Grid 1:
R ____________________
| | | |
| 20 | 21 | 22 |
|______|______|______|
| | | |
| 10 | 11 | 12 |
|______|______|______|
| | | |
| 00 | 01 | 02 |
|______|______|______|
0 R
where 02 means y = 0, x = 2
Now consider this subset of the triangle which corresponds to
grid cells 0 and 8:
..code-block::python
1/R ________
|\ 8 |
| \ |
| 0 \ |
|_______\|
0 1/R
The centroids of the triangles are:
0: (1/3, 1/3) * 1/R
8: (2/3, 2/3) * 1/R
For each grid cell we can now calculate the centroid `(c_y, c_x)`
of the corresponding texture triangle:
- if `(x + y) < R`, then offset the centroid of
triangle 0 by `(y, x) * (1/R)`
- if `(x + y) > R`, then offset the centroid of
triangle 8 by `((R-1-y), (R-1-x)) * (1/R)`.
This is equivalent to updating the portion of Grid 1
above the diagonal, replacing `(y, x)` with `((R-1-y), (R-1-x))`:
..code-block::python
R _____________________
| | | |
| 20 | 01 | 00 |
|______|______|______|
| | | |
| 10 | 11 | 10 |
|______|______|______|
| | | |
| 00 | 01 | 02 |
|______|______|______|
0 R
The barycentric coordinates (w0, w1, w2) are then given by:
..code-block::python
w0 = c_x
w1 = c_y
w2 = 1- w0 - w1
Args:
image: FloatTensor of shape (H, W, 3)
faces_verts_uvs: uv coordinates for each vertex in each face (F, 3, 2)
texture_size: int
Returns:
atlas: a FloatTensor of shape (F, texture_size, texture_size, 3) giving a
per face texture map.
[1] Liu et al, 'Soft Rasterizer: A Differentiable Renderer for Image-based
3D Reasoning', ICCV 2019
"""
R = texture_size
device = faces_verts_uvs.device
rng = torch.arange(R, device=device)
# Meshgrid returns (row, column) i.e (Y, X)
# Change order to (X, Y) to make the grid.
Y, X = meshgrid_ij(rng, rng)
# pyre-fixme[28]: Unexpected keyword argument `axis`.
grid = torch.stack([X, Y], axis=-1) # (R, R, 2)
# Grid cells below the diagonal: x + y < R.
below_diag = grid.sum(-1) < R
# map a [0, R] grid -> to a [0, 1] barycentric coordinates of
# the texture triangle centroids.
bary = torch.zeros((R, R, 3), device=device) # (R, R, 3)
slc = torch.arange(2, device=device)[:, None]
# w0, w1
bary[below_diag, slc] = ((grid[below_diag] + 1.0 / 3.0) / R).T
# w0, w1 for above diagonal grid cells.
# pyre-fixme[16]: `float` has no attribute `T`.
bary[~below_diag,
slc] = (((R - 1.0 - grid[~below_diag]) + 2.0 / 3.0) / R).T
# w2 = 1. - w0 - w1
bary[..., -1] = 1 - bary[..., :2].sum(dim=-1)
# Calculate the uv position in the image for each pixel
# in the per face texture map
# (F, 1, 1, 3, 2) * (R, R, 3, 1) -> (F, R, R, 3, 2) -> (F, R, R, 2)
uv_pos = (faces_verts_uvs[:, None, None] * bary[..., None]).sum(-2)
# bi-linearly interpolate the textures from the images
# using the uv coordinates given by uv_pos.
textures = _bilinear_interpolation_grid_sample(image, uv_pos)
return textures