mirror of
https://github.com/facebookresearch/pytorch3d.git
synced 2025-08-02 20:02:49 +08:00
e12a08133f
Summary: If you join several meshes which have TexturesUV textures using join_meshes_as_scene then we amalgamate all the texture images in to a single one. This now checks if some of the images are equal (i.e. the tensors are the same tensor, in the `is` sense; they have the same `id` in Python) and only uses one copy if they are. I have an example of a massive scene made of several textured meshes with some shared, where this makes the difference between fitting the data on the GPU and not. Reviewed By: theschnitz Differential Revision: D25982364 fbshipit-source-id: a8228805f38475c796302e27328a340d9b56c8ef
316 lines
11 KiB
Python
316 lines
11 KiB
Python
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
|
|
|
|
|
|
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 -------------------- #
|
|
|
|
|
|
class Rectangle(NamedTuple):
|
|
xsize: int
|
|
ysize: int
|
|
identifier: int
|
|
|
|
|
|
class PackedRectangle(NamedTuple):
|
|
x: int
|
|
y: int
|
|
flipped: bool
|
|
is_first: bool
|
|
|
|
|
|
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
|
|
|
|
|
|
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)
|
|
|
|
|
|
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)
|