# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.
# pyre-unsafe
from typing import List, NamedTuple, Tuple
import torch
from pytorch3d.ops import interpolate_face_attributes
def _clip_barycentric_coordinates(bary) -> torch.Tensor:
"""
Args:
bary: barycentric coordinates of shape (...., 3) where `...` represents
an arbitrary number of dimensions
Returns:
bary: Barycentric coordinates clipped (i.e any values < 0 are set to 0)
and renormalized. We only clip the negative values. Values > 1 will fall
into the [0, 1] range after renormalization.
The output is the same shape as the input.
"""
if bary.shape[-1] != 3:
msg = "Expected barycentric coords to have last dim = 3; got %r"
raise ValueError(msg % (bary.shape,))
ndims = bary.ndim - 1
mask = bary.eq(-1).all(dim=-1, keepdim=True).expand(*((-1,) * ndims + (3,)))
clipped = bary.clamp(min=0.0)
clipped[mask] = 0.0
clipped_sum = torch.clamp(clipped.sum(dim=-1, keepdim=True), min=1e-5)
clipped = clipped / clipped_sum
clipped[mask] = -1.0
return clipped
def _interpolate_zbuf(
pix_to_face: torch.Tensor, barycentric_coords: torch.Tensor, meshes
) -> torch.Tensor:
"""
A helper function to calculate the z buffer for each pixel in the
rasterized output.
Args:
pix_to_face: LongTensor of shape (N, H, W, K) specifying the indices
of the faces (in the packed representation) which
overlap each pixel in the image.
barycentric_coords: FloatTensor of shape (N, H, W, K, 3) specifying
the barycentric coordinates of each pixel
relative to the faces (in the packed
representation) which overlap the pixel.
meshes: Meshes object representing a batch of meshes.
Returns:
zbuffer: (N, H, W, K) FloatTensor
"""
verts = meshes.verts_packed()
faces = meshes.faces_packed()
faces_verts_z = verts[faces][..., 2][..., None] # (F, 3, 1)
zbuf = interpolate_face_attributes(pix_to_face, barycentric_coords, faces_verts_z)[
..., 0
] # (1, H, W, K)
zbuf[pix_to_face == -1] = -1
return zbuf
# ----------- Rectangle Packing -------------------- #
[docs]
class Rectangle(NamedTuple):
xsize: int
ysize: int
identifier: int
[docs]
class PackedRectangle(NamedTuple):
x: int
y: int
flipped: bool
is_first: bool
[docs]
class PackedRectangles(NamedTuple):
total_size: Tuple[int, int]
locations: List[PackedRectangle]
# Note the order of members matters here because it determines the queue order.
# We want to place longer rectangles first.
class _UnplacedRectangle(NamedTuple):
size: Tuple[int, int]
ind: int
flipped: bool
def _try_place_rectangle(
rect: _UnplacedRectangle,
placed_so_far: List[PackedRectangle],
occupied: List[Tuple[int, int]],
) -> bool:
"""
Try to place rect within the current bounding box.
Part of the implementation of pack_rectangles.
Note that the arguments `placed_so_far` and `occupied` are modified.
Args:
rect: rectangle to place
placed_so_far: the locations decided upon so far - a list of
(x, y, whether flipped). The nth element is the
location of the nth rectangle if it has been decided.
(modified in place)
occupied: the nodes of the graph of extents of rightmost placed
rectangles - (modified in place)
Returns:
True on success.
Example:
(We always have placed the first rectangle horizontally and other
rectangles above it.)
Let's say the placed boxes 1-4 are laid out like this.
The coordinates of the points marked X are stored in occupied.
It is to the right of the X's that we seek to place rect.
+-----------------------X
|2 |
| +---X
| |4 |
| | |
| +---+X
| |3 |
| | |
+-----------------------+----+------X
y |1 |
^ | --->x |
| +-----------------------------------+
We want to place this rectangle.
+-+
|5|
| |
| | = rect
| |
| |
| |
+-+
The call will succeed, returning True, leaving us with
+-----------------------X
|2 | +-X
| +---+|5|
| |4 || |
| | || |
| +---++ |
| |3 | |
| | | |
+-----------------------+----+-+----X
|1 |
| |
+-----------------------------------+ .
"""
total_width = occupied[0][0]
needed_height = rect.size[1]
current_start_idx = None
current_max_width = 0
previous_height = 0
currently_packed = 0
for idx, interval in enumerate(occupied):
if interval[0] <= total_width - rect.size[0]:
currently_packed += interval[1] - previous_height
current_max_width = max(interval[0], current_max_width)
if current_start_idx is None:
current_start_idx = idx
if currently_packed >= needed_height:
current_max_width = max(interval[0], current_max_width)
placed_so_far[rect.ind] = PackedRectangle(
current_max_width,
occupied[current_start_idx - 1][1],
rect.flipped,
True,
)
new_occupied = (
current_max_width + rect.size[0],
occupied[current_start_idx - 1][1] + needed_height,
)
if currently_packed == needed_height:
occupied[idx] = new_occupied
del occupied[current_start_idx:idx]
elif idx > current_start_idx:
occupied[idx - 1] = new_occupied
del occupied[current_start_idx : (idx - 1)]
else:
occupied.insert(idx, new_occupied)
return True
else:
current_start_idx = None
current_max_width = 0
currently_packed = 0
previous_height = interval[1]
return False
[docs]
def pack_rectangles(sizes: List[Tuple[int, int]]) -> PackedRectangles:
"""
Naive rectangle packing in to a large rectangle. Flipping (i.e. rotating
a rectangle by 90 degrees) is allowed.
This is used to join several uv maps into a single scene, see
TexturesUV.join_scene.
Args:
sizes: List of sizes of rectangles to pack
Returns:
total_size: size of total large rectangle
rectangles: location for each of the input rectangles.
This includes whether they are flipped.
The is_first field is always True.
"""
if len(sizes) < 2:
raise ValueError("Cannot pack less than two boxes")
queue = []
for i, size in enumerate(sizes):
if size[0] < size[1]:
queue.append(_UnplacedRectangle((size[1], size[0]), i, True))
else:
queue.append(_UnplacedRectangle((size[0], size[1]), i, False))
queue.sort()
placed_so_far = [PackedRectangle(-1, -1, False, False)] * len(sizes)
biggest = queue.pop()
total_width, current_height = biggest.size
placed_so_far[biggest.ind] = PackedRectangle(0, 0, biggest.flipped, True)
second = queue.pop()
placed_so_far[second.ind] = PackedRectangle(0, current_height, second.flipped, True)
current_height += second.size[1]
occupied = [biggest.size, (second.size[0], current_height)]
for rect in reversed(queue):
if _try_place_rectangle(rect, placed_so_far, occupied):
continue
rotated = _UnplacedRectangle(
(rect.size[1], rect.size[0]), rect.ind, not rect.flipped
)
if _try_place_rectangle(rotated, placed_so_far, occupied):
continue
# rect wasn't placed in the current bounding box,
# so we add extra space to fit it in.
placed_so_far[rect.ind] = PackedRectangle(0, current_height, rect.flipped, True)
current_height += rect.size[1]
occupied.append((rect.size[0], current_height))
return PackedRectangles((total_width, current_height), placed_so_far)
[docs]
def pack_unique_rectangles(rectangles: List[Rectangle]) -> PackedRectangles:
"""
Naive rectangle packing in to a large rectangle. Flipping (i.e. rotating
a rectangle by 90 degrees) is allowed. Inputs are deduplicated by their
identifier.
This is a wrapper around pack_rectangles, where inputs come with an
identifier. In particular, it calls pack_rectangles for the deduplicated inputs,
then returns the values for all the inputs. The output for all rectangles with
the same identifier will be the same, except that only the first one will have
the is_first field True.
This is used to join several uv maps into a single scene, see
TexturesUV.join_scene.
Args:
rectangles: List of sizes of rectangles to pack
Returns:
total_size: size of total large rectangle
rectangles: location for each of the input rectangles.
This includes whether they are flipped.
The is_first field is true for the first rectangle
with each identifier.
"""
if len(rectangles) < 2:
raise ValueError("Cannot pack less than two boxes")
input_map = {}
input_indices: List[Tuple[int, bool]] = []
unique_input_sizes: List[Tuple[int, int]] = []
for rectangle in rectangles:
if rectangle.identifier not in input_map:
unique_index = len(unique_input_sizes)
unique_input_sizes.append((rectangle.xsize, rectangle.ysize))
input_map[rectangle.identifier] = unique_index
input_indices.append((unique_index, True))
else:
unique_index = input_map[rectangle.identifier]
input_indices.append((unique_index, False))
if len(unique_input_sizes) == 1:
first = [PackedRectangle(0, 0, False, True)]
rest = (len(rectangles) - 1) * [PackedRectangle(0, 0, False, False)]
return PackedRectangles(unique_input_sizes[0], first + rest)
total_size, unique_locations = pack_rectangles(unique_input_sizes)
full_locations = []
for input_index, first in input_indices:
full_locations.append(unique_locations[input_index]._replace(is_first=first))
return PackedRectangles(total_size, full_locations)