mirror of
https://github.com/facebookresearch/pytorch3d.git
synced 2025-12-22 23:30:35 +08:00
Initial commit
fbshipit-source-id: ad58e416e3ceeca85fae0583308968d04e78fe0d
This commit is contained in:
facebook-github-bot
12
pytorch3d/structures/__init__.py
Normal file
12
pytorch3d/structures/__init__.py
Normal file
@@ -0,0 +1,12 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
|
||||
|
||||
from .meshes import Meshes
|
||||
from .textures import Textures
|
||||
from .utils import (
|
||||
list_to_packed,
|
||||
list_to_padded,
|
||||
packed_to_list,
|
||||
padded_to_list,
|
||||
)
|
||||
|
||||
__all__ = [k for k in globals().keys() if not k.startswith("_")]
|
||||
1367
pytorch3d/structures/meshes.py
Normal file
1367
pytorch3d/structures/meshes.py
Normal file
File diff suppressed because it is too large
Load Diff
205
pytorch3d/structures/textures.py
Normal file
205
pytorch3d/structures/textures.py
Normal file
@@ -0,0 +1,205 @@
|
||||
#!/usr/bin/env python3
|
||||
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
|
||||
|
||||
from typing import List, Union
|
||||
import torch
|
||||
import torchvision.transforms as T
|
||||
|
||||
from .utils import list_to_packed, padded_to_list
|
||||
|
||||
|
||||
"""
|
||||
This file has functions for interpolating textures after rasterization.
|
||||
"""
|
||||
|
||||
|
||||
def _pad_texture_maps(images: List[torch.Tensor]) -> torch.Tensor:
|
||||
"""
|
||||
Pad all texture images so they have the same height and width.
|
||||
|
||||
Args:
|
||||
images: list of N tensors of shape (H, W)
|
||||
|
||||
Returns:
|
||||
tex_maps: Tensor of shape (N, max_H, max_W)
|
||||
"""
|
||||
tex_maps = []
|
||||
max_H = 0
|
||||
max_W = 0
|
||||
for im in images:
|
||||
h, w, _3 = im.shape
|
||||
if h > max_H:
|
||||
max_H = h
|
||||
if w > max_W:
|
||||
max_W = w
|
||||
tex_maps.append(im)
|
||||
max_shape = (max_H, max_W)
|
||||
|
||||
# If all texture images are not the same size then resize to the
|
||||
# largest size.
|
||||
resize = T.Compose([T.ToPILImage(), T.Resize(size=max_shape), T.ToTensor()])
|
||||
|
||||
for i, image in enumerate(tex_maps):
|
||||
if image.shape != max_shape:
|
||||
# ToPIL takes and returns a C x H x W tensor
|
||||
image = resize(image.permute(2, 0, 1)).permute(1, 2, 0)
|
||||
tex_maps[i] = image
|
||||
tex_maps = torch.stack(tex_maps, dim=0) # (num_tex_maps, max_H, max_W, 3)
|
||||
return tex_maps
|
||||
|
||||
|
||||
def _extend_tensor(input_tensor: torch.Tensor, N: int) -> torch.Tensor:
|
||||
"""
|
||||
Extend a tensor `input_tensor` with ndim > 2, `N` times along the batch
|
||||
dimension. This is done in the following sequence of steps (where `B` is
|
||||
the batch dimension):
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
input_tensor (B, ...)
|
||||
-> add leading empty dimension (1, B, ...)
|
||||
-> expand (N, B, ...)
|
||||
-> reshape (N * B, ...)
|
||||
|
||||
Args:
|
||||
input_tensor: torch.Tensor with ndim > 2 representing a batched input.
|
||||
N: number of times to extend each element of the batch.
|
||||
"""
|
||||
if input_tensor.ndim < 2:
|
||||
raise ValueError("Input tensor must have ndimensions >= 2.")
|
||||
B = input_tensor.shape[0]
|
||||
non_batch_dims = tuple(input_tensor.shape[1:])
|
||||
constant_dims = (-1,) * input_tensor.ndim # these dims are not expanded.
|
||||
return (
|
||||
input_tensor.clone()[None, ...]
|
||||
.expand(N, *constant_dims)
|
||||
.transpose(0, 1)
|
||||
.reshape(N * B, *non_batch_dims)
|
||||
)
|
||||
|
||||
|
||||
class Textures(object):
|
||||
def __init__(
|
||||
self,
|
||||
maps: Union[List, torch.Tensor] = None,
|
||||
faces_uvs: torch.Tensor = None,
|
||||
verts_uvs: torch.Tensor = None,
|
||||
verts_rgb: torch.Tensor = None,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
maps: texture map per mesh. This can either be a list of maps
|
||||
[(H, W, 3)] or a padded tensor of shape (N, H, W, 3).
|
||||
faces_uvs: (N, F, 3) tensor giving the index into verts_uvs for each
|
||||
vertex in the face. Padding value is assumed to be -1.
|
||||
verts_uvs: (N, V, 2) tensor giving the uv coordinate per vertex.
|
||||
verts_rgb: (N, V, 3) tensor giving the rgb color per vertex.
|
||||
"""
|
||||
if faces_uvs is not None and faces_uvs.ndim != 3:
|
||||
msg = "Expected faces_uvs to be of shape (N, F, 3); got %r"
|
||||
raise ValueError(msg % repr(faces_uvs.shape))
|
||||
if verts_uvs is not None and verts_uvs.ndim != 3:
|
||||
msg = "Expected verts_uvs to be of shape (N, V, 2); got %r"
|
||||
raise ValueError(msg % repr(faces_uvs.shape))
|
||||
if verts_rgb is not None and verts_rgb.ndim != 3:
|
||||
msg = "Expected verts_rgb to be of shape (N, V, 3); got %r"
|
||||
raise ValueError(msg % verts_rgb.shape)
|
||||
if maps is not None:
|
||||
if torch.is_tensor(map) and map.ndim != 4:
|
||||
msg = "Expected maps to be of shape (N, H, W, 3); got %r"
|
||||
raise ValueError(msg % repr(maps.shape))
|
||||
elif isinstance(maps, list):
|
||||
maps = _pad_texture_maps(maps)
|
||||
self._faces_uvs_padded = faces_uvs
|
||||
self._verts_uvs_padded = verts_uvs
|
||||
self._verts_rgb_padded = verts_rgb
|
||||
self._maps_padded = maps
|
||||
self._num_faces_per_mesh = None
|
||||
|
||||
if self._faces_uvs_padded is not None:
|
||||
self._num_faces_per_mesh = faces_uvs.gt(-1).all(-1).sum(-1).tolist()
|
||||
|
||||
def clone(self):
|
||||
other = Textures()
|
||||
for k in dir(self):
|
||||
v = getattr(self, k)
|
||||
if torch.is_tensor(v):
|
||||
setattr(other, k, v.clone())
|
||||
return other
|
||||
|
||||
def to(self, device):
|
||||
for k in dir(self):
|
||||
v = getattr(self, k)
|
||||
if torch.is_tensor(v) and v.device != device:
|
||||
setattr(self, k, v.to(device))
|
||||
return self
|
||||
|
||||
def faces_uvs_padded(self) -> torch.Tensor:
|
||||
return self._faces_uvs_padded
|
||||
|
||||
def faces_uvs_list(self) -> List[torch.Tensor]:
|
||||
if self._faces_uvs_padded is not None:
|
||||
return padded_to_list(
|
||||
self._faces_uvs_padded, split_size=self._num_faces_per_mesh
|
||||
)
|
||||
|
||||
def faces_uvs_packed(self) -> torch.Tensor:
|
||||
return list_to_packed(self.faces_uvs_list())[0]
|
||||
|
||||
def verts_uvs_padded(self) -> torch.Tensor:
|
||||
return self._verts_uvs_padded
|
||||
|
||||
def verts_uvs_list(self) -> List[torch.Tensor]:
|
||||
return padded_to_list(self._verts_uvs_padded)
|
||||
|
||||
def verts_uvs_packed(self) -> torch.Tensor:
|
||||
return list_to_packed(self.verts_uvs_list())[0]
|
||||
|
||||
def verts_rgb_padded(self) -> torch.Tensor:
|
||||
return self._verts_rgb_padded
|
||||
|
||||
def verts_rgb_list(self) -> List[torch.Tensor]:
|
||||
return padded_to_list(self._verts_rgb_padded)
|
||||
|
||||
def verts_rgb_packed(self) -> torch.Tensor:
|
||||
return list_to_packed(self.verts_rgb_list())[0]
|
||||
|
||||
# Currently only the padded maps are used.
|
||||
def maps_padded(self) -> torch.Tensor:
|
||||
return self._maps_padded
|
||||
|
||||
def extend(self, N: int) -> "Textures":
|
||||
"""
|
||||
Create new Textures class which contains each input texture N times
|
||||
|
||||
Args:
|
||||
N: number of new copies of each texture.
|
||||
|
||||
Returns:
|
||||
new Textures object.
|
||||
"""
|
||||
if not isinstance(N, int):
|
||||
raise ValueError("N must be an integer.")
|
||||
if N <= 0:
|
||||
raise ValueError("N must be > 0.")
|
||||
|
||||
if all(
|
||||
v is not None
|
||||
for v in [
|
||||
self._faces_uvs_padded,
|
||||
self._verts_uvs_padded,
|
||||
self._maps_padded,
|
||||
]
|
||||
):
|
||||
new_verts_uvs = _extend_tensor(self._verts_uvs_padded, N)
|
||||
new_faces_uvs = _extend_tensor(self._faces_uvs_padded, N)
|
||||
new_maps = _extend_tensor(self._maps_padded, N)
|
||||
return Textures(
|
||||
verts_uvs=new_verts_uvs, faces_uvs=new_faces_uvs, maps=new_maps
|
||||
)
|
||||
elif self._verts_rgb_padded is not None:
|
||||
new_verts_rgb = _extend_tensor(self._verts_rgb_padded, N)
|
||||
return Textures(verts_rgb=new_verts_rgb)
|
||||
else:
|
||||
msg = "Either vertex colors or texture maps are required."
|
||||
raise ValueError(msg)
|
||||
150
pytorch3d/structures/utils.py
Normal file
150
pytorch3d/structures/utils.py
Normal file
@@ -0,0 +1,150 @@
|
||||
#!/usr/bin/env python3
|
||||
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
|
||||
|
||||
from typing import List, Union
|
||||
import torch
|
||||
|
||||
|
||||
"""
|
||||
Util functions containing representation transforms for points/verts/faces.
|
||||
"""
|
||||
|
||||
|
||||
def list_to_padded(
|
||||
x: List[torch.Tensor],
|
||||
pad_size: Union[list, tuple, None] = None,
|
||||
pad_value: float = 0.0,
|
||||
equisized: bool = False,
|
||||
) -> torch.Tensor:
|
||||
r"""
|
||||
Transforms a list of N tensors each of shape (Mi, Ki) into a single tensor
|
||||
of shape (N, pad_size(0), pad_size(1)), or (N, max(Mi), max(Ki))
|
||||
if pad_size is None.
|
||||
|
||||
Args:
|
||||
x: list of Tensors
|
||||
pad_size: list(int) specifying the size of the padded tensor
|
||||
pad_value: float value to be used to fill the padded tensor
|
||||
equisized: bool indicating whether the items in x are of equal size
|
||||
(sometimes this is known and if provided saves computation)
|
||||
|
||||
Returns:
|
||||
x_padded: tensor consisting of padded input tensors
|
||||
"""
|
||||
if equisized:
|
||||
return torch.stack(x, 0)
|
||||
|
||||
if pad_size is None:
|
||||
pad_dim0 = max(y.shape[0] for y in x if len(y) > 0)
|
||||
pad_dim1 = max(y.shape[1] for y in x if len(y) > 0)
|
||||
else:
|
||||
if len(pad_size) != 2:
|
||||
raise ValueError(
|
||||
"Pad size must contain target size for 1st and 2nd dim"
|
||||
)
|
||||
pad_dim0, pad_dim1 = pad_size
|
||||
|
||||
N = len(x)
|
||||
x_padded = torch.full(
|
||||
(N, pad_dim0, pad_dim1), pad_value, dtype=x[0].dtype, device=x[0].device
|
||||
)
|
||||
for i, y in enumerate(x):
|
||||
if len(y) > 0:
|
||||
if y.ndim != 2:
|
||||
raise ValueError("Supports only 2-dimensional tensor items")
|
||||
x_padded[i, : y.shape[0], : y.shape[1]] = y
|
||||
return x_padded
|
||||
|
||||
|
||||
def padded_to_list(
|
||||
x: torch.Tensor, split_size: Union[list, tuple, None] = None
|
||||
):
|
||||
r"""
|
||||
Transforms a padded tensor of shape (N, M, K) into a list of N tensors
|
||||
of shape (Mi, Ki) where (Mi, Ki) is specified in split_size(i), or of shape
|
||||
(M, K) if split_size is None.
|
||||
Support only for 3-dimensional input tensor.
|
||||
|
||||
Args:
|
||||
x: tensor
|
||||
split_size: the shape of the final tensor to be returned (of length N).
|
||||
"""
|
||||
if x.ndim != 3:
|
||||
raise ValueError("Supports only 3-dimensional input tensors")
|
||||
x_list = list(x.unbind(0))
|
||||
if split_size is None:
|
||||
return x_list
|
||||
|
||||
N = len(split_size)
|
||||
if x.shape[0] != N:
|
||||
raise ValueError(
|
||||
"Split size must be of same length as inputs first dimension"
|
||||
)
|
||||
|
||||
for i in range(N):
|
||||
if isinstance(split_size[i], int):
|
||||
x_list[i] = x_list[i][: split_size[i]]
|
||||
elif len(split_size[i]) == 2:
|
||||
x_list[i] = x_list[i][: split_size[i][0], : split_size[i][1]]
|
||||
else:
|
||||
raise ValueError(
|
||||
"Support only for 2-dimensional unbinded tensor. \
|
||||
Split size for more dimensions provided"
|
||||
)
|
||||
return x_list
|
||||
|
||||
|
||||
def list_to_packed(x: List[torch.Tensor]):
|
||||
r"""
|
||||
Transforms a list of N tensors each of shape (Mi, K, ...) into a single
|
||||
tensor of shape (sum(Mi), K, ...).
|
||||
|
||||
Args:
|
||||
x: list of tensors.
|
||||
|
||||
Returns:
|
||||
4-element tuple containing
|
||||
|
||||
- **x_packed**: tensor consisting of packed input tensors along the
|
||||
1st dimension.
|
||||
- **num_items**: tensor of shape N containing Mi for each element in x.
|
||||
- **item_packed_first_idx**: tensor of shape N indicating the index of
|
||||
the first item belonging to the same element in the original list.
|
||||
- **item_packed_to_list_idx**: tensor of shape sum(Mi) containing the
|
||||
index of the element in the list the item belongs to.
|
||||
"""
|
||||
N = len(x)
|
||||
num_items = torch.zeros(N, dtype=torch.int64, device=x[0].device)
|
||||
item_packed_first_idx = torch.zeros(
|
||||
N, dtype=torch.int64, device=x[0].device
|
||||
)
|
||||
item_packed_to_list_idx = []
|
||||
cur = 0
|
||||
for i, y in enumerate(x):
|
||||
num = len(y)
|
||||
num_items[i] = num
|
||||
item_packed_first_idx[i] = cur
|
||||
item_packed_to_list_idx.append(
|
||||
torch.full((num,), i, dtype=torch.int64, device=y.device)
|
||||
)
|
||||
cur += num
|
||||
|
||||
x_packed = torch.cat(x, dim=0)
|
||||
item_packed_to_list_idx = torch.cat(item_packed_to_list_idx, dim=0)
|
||||
|
||||
return x_packed, num_items, item_packed_first_idx, item_packed_to_list_idx
|
||||
|
||||
|
||||
def packed_to_list(x: torch.Tensor, split_size: Union[list, int]):
|
||||
r"""
|
||||
Transforms a tensor of shape (sum(Mi), K, L, ...) to N set of tensors of
|
||||
shape (Mi, K, L, ...) where Mi's are defined in split_size
|
||||
|
||||
Args:
|
||||
x: tensor
|
||||
split_size: list or int defining the number of items for each split
|
||||
|
||||
Returns:
|
||||
x_list: A list of Tensors
|
||||
"""
|
||||
return x.split(split_size, dim=0)
|
||||
Reference in New Issue
Block a user