Skip to content

GitLab

Commit dbf06b50 authored by facebook-github-bot's avatar facebook-github-bot
Browse files

Initial commit 

fbshipit-source-id: ad58e416e3ceeca85fae0583308968d04e78fe0d
parents
Hide whitespace changes
Inline Side-by-side
Showing with 3868 additions and 0 deletions
pytorch3d/renderer/mesh/texturing.py 0 → 100644
View file @ dbf06b50
#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
import torch
import torch.nn.functional as F
from pytorch3d.structures.textures import Textures
def _clip_barycentric_coordinates(bary) -> torch.Tensor:
"""
Args:
bary: barycentric coordinates of shape (...., 3) where `...` represents
an arbitrary number of dimensions
Returns:
bary: All barycentric coordinate values clipped to the range [0, 1]
and renormalized. 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)
clipped = bary.clamp(min=0, max=1)
clipped_sum = torch.clamp(clipped.sum(dim=-1, keepdim=True), min=1e-5)
clipped = clipped / clipped_sum
return clipped
def interpolate_face_attributes(
fragments, face_attributes: torch.Tensor, bary_clip: bool = False
) -> torch.Tensor:
"""
Interpolate arbitrary face attributes using the barycentric coordinates
for each pixel in the rasterized output.
Args:
fragments:
The outputs of rasterization. From this we use
- 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 coordianates of each pixel
relative to the faces (in the packed
representation) which overlap the pixel.
face_attributes: packed attributes of shape (total_faces, 3, D),
specifying the value of the attribute for each
vertex in the face.
bary_clip: Bool to indicate if barycentric_coords should be clipped
before being used for interpolation.
Returns:
pixel_vals: tensor of shape (N, H, W, K, D) giving the interpolated
value of the face attribute for each pixel.
"""
pix_to_face = fragments.pix_to_face
barycentric_coords = fragments.bary_coords
F, FV, D = face_attributes.shape
if FV != 3:
raise ValueError("Faces can only have three vertices; got %r" % FV)
N, H, W, K, _ = barycentric_coords.shape
if pix_to_face.shape != (N, H, W, K):
msg = "pix_to_face must have shape (batch_size, H, W, K); got %r"
raise ValueError(msg % pix_to_face.shape)
if bary_clip:
barycentric_coords = _clip_barycentric_coordinates(barycentric_coords)
# Replace empty pixels in pix_to_face with 0 in order to interpolate.
mask = pix_to_face == -1
pix_to_face = pix_to_face.clone()
pix_to_face[mask] = 0
idx = pix_to_face.view(N * H * W * K, 1, 1).expand(N * H * W * K, 3, D)
pixel_face_vals = face_attributes.gather(0, idx).view(N, H, W, K, 3, D)
pixel_vals = (barycentric_coords[..., None] * pixel_face_vals).sum(dim=-2)
pixel_vals[mask] = 0 # Replace masked values in output.
return pixel_vals
def interpolate_texture_map(fragments, meshes) -> torch.Tensor:
"""
Interpolate a 2D texture map using uv vertex texture coordinates for each
face in the mesh. First interpolate the vertex uvs using barycentric coordinates
for each pixel in the rasterized output. Then interpolate the texture map
using the uv coordinate for each pixel.
Args:
fragments:
The outputs of rasterization. From this we use
- 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 coordianates of each pixel
relative to the faces (in the packed
representation) which overlap the pixel.
meshes: Meshes representing a batch of meshes. It is expected that
meshes has a textures attribute which is an instance of the
Textures class.
Returns:
texels: tensor of shape (N, H, W, K, C) giving the interpolated
texture for each pixel in the rasterized image.
"""
if not isinstance(meshes.textures, Textures):
msg = "Expected meshes.textures to be an instance of Textures; got %r"
raise ValueError(msg % type(meshes.textures))
faces_uvs = meshes.textures.faces_uvs_packed()
verts_uvs = meshes.textures.verts_uvs_packed()
faces_verts_uvs = verts_uvs[faces_uvs]
texture_maps = meshes.textures.maps_padded()
# pixel_uvs: (N, H, W, K, 2)
pixel_uvs = interpolate_face_attributes(fragments, faces_verts_uvs)
N, H_out, W_out, K = fragments.pix_to_face.shape
N, H_in, W_in, C = texture_maps.shape # 3 for RGB
# pixel_uvs: (N, H, W, K, 2) -> (N, K, H, W, 2) -> (NK, H, W, 2)
pixel_uvs = pixel_uvs.permute(0, 3, 1, 2, 4).view(N * K, H_out, W_out, 2)
# textures.map:
# (N, H, W, C) -> (N, C, H, W) -> (1, N, C, H, W)
# -> expand (K, N, C, H, W) -> reshape (N*K, C, H, W)
texture_maps = (
texture_maps.permute(0, 3, 1, 2)[None, ...]
.expand(K, -1, -1, -1, -1)
.transpose(0, 1)
.reshape(N * K, C, H_in, W_in)
)
# Textures: (N*K, C, H, W), pixel_uvs: (N*K, H, W, 2)
# Now need to format the pixel uvs and the texture map correctly!
# From pytorch docs, grid_sample takes `grid` and `input`:
# grid specifies the sampling pixel locations normalized by
# the input spatial dimensions It should have most
# values in the range of [-1, 1]. Values x = -1, y = -1
# is the left-top pixel of input, and values x = 1, y = 1 is the
# right-bottom pixel of input.
pixel_uvs = pixel_uvs * 2.0 - 1.0
texture_maps = torch.flip(
texture_maps, [2]
) # flip y axis of the texture map
if texture_maps.device != pixel_uvs.device:
texture_maps = texture_maps.to(pixel_uvs.device)
texels = F.grid_sample(texture_maps, pixel_uvs, align_corners=False)
texels = texels.view(N, K, C, H_out, W_out).permute(0, 3, 4, 1, 2)
return texels
def interpolate_vertex_colors(fragments, meshes) -> torch.Tensor:
"""
Detemine the color for each rasterized face. Interpolate the colors for
vertices which form the face using the barycentric coordinates.
Args:
meshes: A Meshes class representing a batch of meshes.
fragments:
The outputs of rasterization. From this we use
- 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 coordianates of each pixel
relative to the faces (in the packed
representation) which overlap the pixel.
Returns:
texels: An texture per pixel of shape (N, H, W, K, C).
There will be one C dimensional value for each element in
fragments.pix_to_face.
"""
vertex_textures = meshes.textures.verts_rgb_padded().view(-1, 3) # (V, C)
vertex_textures = vertex_textures[meshes.verts_padded_to_packed_idx(), :]
faces_packed = meshes.faces_packed()
faces_textures = vertex_textures[faces_packed] # (F, 3, C)
texels = interpolate_face_attributes(fragments, faces_textures)
return texels
pytorch3d/renderer/utils.py 0 → 100644
View file @ dbf06b50
#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
import numpy as np
from typing import Any, Union
import torch
class TensorAccessor(object):
"""
A helper class to be used with the __getitem__ method. This can be used for
getting/setting the values for an attribute of a class at one particular
index. This is useful when the attributes of a class are batched tensors
and one element in the batch needs to be modified.
"""
def __init__(self, class_object, index: Union[int, slice]):
"""
Args:
class_object: this should be an instance of a class which has
attributes which are tensors representing a batch of
values.
index: int/slice, an index indicating the position in the batch.
In __setattr__ and __getattr__ only the value of class
attributes at this index will be accessed.
"""
self.__dict__["class_object"] = class_object
self.__dict__["index"] = index
def __setattr__(self, name: str, value: Any):
"""
Update the attribute given by `name` to the value given by `value`
at the index specified by `self.index`.
Args:
name: str, name of the attribute.
value: value to set the attribute to.
"""
v = getattr(self.class_object, name)
if not torch.is_tensor(v):
msg = "Can only set values on attributes which are tensors; got %r"
raise AttributeError(msg % type(v))
# Convert the attribute to a tensor if it is not a tensor.
if not torch.is_tensor(value):
value = torch.tensor(
value,
device=v.device,
dtype=v.dtype,
requires_grad=v.requires_grad,
)
# Check the shapes match the existing shape and the shape of the index.
if v.dim() > 1 and value.dim() > 1 and value.shape[1:] != v.shape[1:]:
msg = "Expected value to have shape %r; got %r"
raise ValueError(msg % (v.shape, value.shape))
if (
v.dim() == 0
and isinstance(self.index, slice)
and len(value) != len(self.index)
):
msg = "Expected value to have len %r; got %r"
raise ValueError(msg % (len(self.index), len(value)))
self.class_object.__dict__[name][self.index] = value
def __getattr__(self, name: str):
"""
Return the value of the attribute given by "name" on self.class_object
at the index specified in self.index.
Args:
name: string of the attribute name
"""
if hasattr(self.class_object, name):
return self.class_object.__dict__[name][self.index]
else:
msg = "Attribue %s not found on %r"
return AttributeError(msg % (name, self.class_object.__name__))
BROADCAST_TYPES = (float, int, list, tuple, torch.Tensor, np.ndarray)
class TensorProperties(object):
"""
A mix-in class for storing tensors as properties with helper methods.
"""
def __init__(self, dtype=torch.float32, device="cpu", **kwargs):
"""
Args:
dtype: data type to set for the inputs
device: str or torch.device
kwargs: any number of keyword arguments. Any arguments which are
of type (float/int/tuple/tensor/array) are broadcasted and
other keyword arguments are set as attributes.
"""
super().__init__()
self.device = device
self._N = 0
if kwargs is not None:
# broadcast all inputs which are float/int/list/tuple/tensor/array
# set as attributes anything else e.g. strings, bools
args_to_broadcast = {}
for k, v in kwargs.items():
if isinstance(v, (str, bool)):
setattr(self, k, v)
elif isinstance(v, BROADCAST_TYPES):
args_to_broadcast[k] = v
else:
msg = "Arg %s with type %r is not broadcastable"
print(msg % (k, type(v)))
names = args_to_broadcast.keys()
# convert from type dict.values to tuple
values = tuple(v for v in args_to_broadcast.values())
if len(values) > 0:
broadcasted_values = convert_to_tensors_and_broadcast(
*values, device=device
)
# Set broadcasted values as attributes on self.
for i, n in enumerate(names):
setattr(self, n, broadcasted_values[i])
if self._N == 0:
self._N = broadcasted_values[i].shape[0]
def __len__(self) -> int:
return self._N
def isempty(self) -> bool:
return self._N == 0
def __getitem__(self, index: Union[int, slice]):
"""
Args:
index: an int or slice used to index all the fields.
Returns:
if `index` is an index int/slice return a TensorAccessor class
with getattribute/setattribute methods which return/update the value
at the index in the original camera.
"""
if isinstance(index, (int, slice)):
return TensorAccessor(class_object=self, index=index)
msg = "Expected index of type int or slice; got %r"
raise ValueError(msg % type(index))
def to(self, device: str = "cpu"):
"""
In place operation to move class properties which are tensors to a
specified device. If self has a property "device", update this as well.
"""
for k in dir(self):
v = getattr(self, k)
if k == "device":
setattr(self, k, device)
if torch.is_tensor(v) and v.device != device:
setattr(self, k, v.to(device))
return self
def clone(self, other):
"""
Update the tensor properties of other with the cloned properties of self.
"""
for k in dir(self):
v = getattr(self, k)
if k == "device":
setattr(self, k, v)
if torch.is_tensor(v):
setattr(other, k, v.clone())
return other
def gather_props(self, batch_idx):
"""
This is an in place operation to reformat all tensor class attributes
based on a set of given indices using torch.gather. This is useful when
attributes which are batched tensors e.g. shape (N, 3) need to be
multiplied with another tensor which has a different first dimension
e.g. packed vertices of shape (V, 3).
Example
.. code-block:: python
self.specular_color = (N, 3) tensor of specular colors for each mesh
A lighting calculation may use
.. code-block:: python
verts_packed = meshes.verts_packed() # (V, 3)
To multiply these two tensors the batch dimension needs to be the same.
To achieve this we can do
.. code-block:: python
batch_idx = meshes.verts_packed_to_mesh_idx() # (V)
This gives index of the mesh for each vertex in verts_packed.
.. code-block:: python
self.gather_props(batch_idx)
self.specular_color = (V, 3) tensor with the specular color for
each packed vertex.
torch.gather requires the index tensor to have the same shape as the
input tensor so this method takes care of the reshaping of the index
tensor to use with class attributes with arbitrary dimensions.
Args:
batch_idx: shape (B, ...) where `...` represents an arbitrary
number of dimensions
Returns:
self with all properties reshaped. e.g. a property with shape (N, 3)
is transformed to shape (B, 3).
"""
for k in dir(self):
v = getattr(self, k)
if torch.is_tensor(v):
if v.shape[0] > 1:
# There are different values for each batch element
# so gather these using the batch_idx
idx_dims = batch_idx.shape
tensor_dims = v.shape
if len(idx_dims) > len(tensor_dims):
msg = "batch_idx cannot have more dimensions than %s. "
msg += "got shape %r and %s has shape %r"
raise ValueError(msg % (k, idx_dims, k, tensor_dims))
if idx_dims != tensor_dims:
# To use torch.gather the index tensor (batch_idx) has
# to have the same shape as the input tensor.
new_dims = len(tensor_dims) - len(idx_dims)
new_shape = idx_dims + (1,) * new_dims
expand_dims = (-1,) + tensor_dims[1:]
batch_idx = batch_idx.view(*new_shape)
batch_idx = batch_idx.expand(*expand_dims)
v = v.gather(0, batch_idx)
setattr(self, k, v)
return self
def format_tensor(
input, dtype=torch.float32, device: str = "cpu"
) -> torch.Tensor:
"""
Helper function for converting a scalar value to a tensor.
Args:
input: Python scalar, Python list/tuple, torch scalar, 1D torch tensor
dtype: data type for the input
device: torch device on which the tensor should be placed.
Returns:
input_vec: torch tensor with optional added batch dimension.
"""
if not torch.is_tensor(input):
input = torch.tensor(input, dtype=dtype, device=device)
if input.dim() == 0:
input = input.view(1)
if input.device != device:
input = input.to(device=device)
return input
def convert_to_tensors_and_broadcast(
*args, dtype=torch.float32, device: str = "cpu"
):
"""
Helper function to handle parsing an arbitrary number of inputs (*args)
which all need to have the same batch dimension.
The output is a list of tensors.
Args:
*args: an arbitrary number of inputs
Each of the values in `args` can be one of the following
- Python scalar
- Torch scalar
- Torch tensor of shape (N, K_i) or (1, K_i) where K_i are
an arbitrary number of dimensions which can vary for each
value in args. In this case each input is broadcast to a
tensor of shape (N, K_i)
dtype: data type to use when creating new tensors.
device: torch device on which the tensors should be placed.
Output:
args: A list of tensors of shape (N, K_i)
"""
# Convert all inputs to tensors with a batch dimension
args_1d = [format_tensor(c, dtype, device) for c in args]
# Find broadcast size
sizes = [c.shape[0] for c in args_1d]
N = max(sizes)
args_Nd = []
for c in args_1d:
if c.shape[0] != 1 and c.shape[0] != N:
msg = "Got non-broadcastable sizes %r" % (sizes)
raise ValueError(msg)
# Expand broadcast dim and keep non broadcast dims the same size
expand_sizes = (N,) + (-1,) * len(c.shape[1:])
args_Nd.append(c.expand(*expand_sizes))
if len(args) == 1:
args_Nd = args_Nd[0] # Return the first element
return args_Nd
pytorch3d/structures/__init__.py 0 → 100644
View file @ dbf06b50
# 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("_")]
pytorch3d/structures/meshes.py 0 → 100644
View file @ dbf06b50
#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
import torch
from pytorch3d import _C
from . import utils as struct_utils
from .textures import Textures
class Meshes(object):
"""
This class provides functions for working with batches of triangulated
meshes with varying numbers of faces and vertices, and converting between
representations.
Within Meshes, there are three different representations of the faces and
verts data:
List
- only used for input as a starting point to convert to other representations.
Padded
- has specific batch dimension.
Packed
- no batch dimension.
- has auxillary variables used to index into the padded representation.
Example:
Input list of verts V_n = [[V_1], [V_2], ... , [V_N]]
where V_1, ... , V_N are the number of verts in each mesh and N is the
numer of meshes.
Input list of faces F_n = [[F_1], [F_2], ... , [F_N]]
where F_1, ... , F_N are the number of faces in each mesh.
# SPHINX IGNORE
List | Padded | Packed
---------------------------|-------------------------|------------------------
[[V_1], ... , [V_N]] | size = (N, max(V_n), 3) | size = (sum(V_n), 3)
| |
Example for verts: | |
| |
V_1 = 3, V_2 = 4, V_3 = 5 | size = (3, 5, 3) | size = (12, 3)
| |
List([ | tensor([ | tensor([
[ | [ | [0.1, 0.3, 0.5],
[0.1, 0.3, 0.5], | [0.1, 0.3, 0.5], | [0.5, 0.2, 0.1],
[0.5, 0.2, 0.1], | [0.5, 0.2, 0.1], | [0.6, 0.8, 0.7],
[0.6, 0.8, 0.7], | [0.6, 0.8, 0.7], | [0.1, 0.3, 0.3],
], | [0, 0, 0], | [0.6, 0.7, 0.8],
[ | [0, 0, 0], | [0.2, 0.3, 0.4],
[0.1, 0.3, 0.3], | ], | [0.1, 0.5, 0.3],
[0.6, 0.7, 0.8], | [ | [0.7, 0.3, 0.6],
[0.2, 0.3, 0.4], | [0.1, 0.3, 0.3], | [0.2, 0.4, 0.8],
[0.1, 0.5, 0.3], | [0.6, 0.7, 0.8], | [0.9, 0.5, 0.2],
], | [0.2, 0.3, 0.4], | [0.2, 0.3, 0.4],
[ | [0.1, 0.5, 0.3], | [0.9, 0.3, 0.8],
[0.7, 0.3, 0.6], | [0, 0, 0], | ])
[0.2, 0.4, 0.8], | ], |
[0.9, 0.5, 0.2], | [ |
[0.2, 0.3, 0.4], | [0.7, 0.3, 0.6], |
[0.9, 0.3, 0.8], | [0.2, 0.4, 0.8], |
] | [0.9, 0.5, 0.2], |
]) | [0.2, 0.3, 0.4], |
| [0.9, 0.3, 0.8], |
| ] |
| ]) |
Example for faces: | |
| |
F_1 = 1, F_2 = 2, F_3 = 7 | size = (3, 7, 3) | size = (10, 3)
| |
List([ | tensor([ | tensor([
[ | [ | [ 0, 1, 2],
[0, 1, 2], | [0, 1, 2], | [ 3, 4, 5],
], | [-1, -1, -1], | [ 4, 5, 6],
[ | [-1, -1, -1] | [ 8, 9, 7],
[0, 1, 2], | [-1, -1, -1] | [ 7, 8, 10],
[1, 2, 3], | [-1, -1, -1] | [ 9, 10, 8],
], | [-1, -1, -1], | [11, 10, 9],
[ | [-1, -1, -1], | [11, 7, 8],
[1, 2, 0], | ], | [11, 10, 8],
[0, 1, 3], | [ | [11, 9, 8],
[2, 3, 1], | [0, 1, 2], | ])
[4, 3, 2], | [1, 2, 3], |
[4, 0, 1], | [-1, -1, -1], |
[4, 3, 1], | [-1, -1, -1], |
[4, 2, 1], | [-1, -1, -1], |
], | [-1, -1, -1], |
]) | [-1, -1, -1], |
| ], |
| [ |
| [1, 2, 0], |
| [0, 1, 3], |
| [2, 3, 1], |
| [4, 3, 2], |
| [4, 0, 1], |
| [4, 3, 1], |
| [4, 2, 1], |
| ] |
| ]) |
-----------------------------------------------------------------------------
Auxillary variables for packed representation
Name | Size | Example from above
-------------------------------|---------------------|-----------------------
| |
verts_packed_to_mesh_idx | size = (sum(V_n)) | tensor([
| | 0, 0, 0, 1, 1, 1,
| | 1, 2, 2, 2, 2, 2
| | )]
| | size = (12)
| |
mesh_to_verts_packed_first_idx | size = (N) | tensor([0, 3, 7])
| | size = (3)
| |
num_verts_per_mesh | size = (N) | tensor([3, 4, 5])
| | size = (3)
| |
faces_packed_to_mesh_idx | size = (sum(F_n)) | tensor([
| | 0, 1, 1, 2, 2, 2,
| | 2, 2, 2, 2
| | )]
| | size = (10)
| |
mesh_to_faces_packed_first_idx | size = (N) | tensor([0, 1, 3])
| | size = (3)
| |
num_faces_per_mesh | size = (N) | tensor([1, 2, 7])
| | size = (3)
| |
verts_padded_to_packed_idx | size = (sum(V_n)) | tensor([
| | 0, 1, 2, 5, 6, 7,
| | 8, 10, 11, 12, 13,
| | 14
| | )]
| | size = (12)
-----------------------------------------------------------------------------
# SPHINX IGNORE
From the faces, edges are computed and have packed and padded
representations with auxillary variables.
E_n = [[E_1], ... , [E_N]]
where E_1, ... , E_N are the number of unique edges in each mesh.
Total number of unique edges = sum(E_n)
# SPHINX IGNORE
Name | Size | Example from above
------------------------------|-------------------------|----------------------
| |
edges_packed | size = (sum(E_n), 2) | tensor([
| | [0, 1],
| | [0, 2],
| | [1, 2],
| | ...
| | [10, 11],
| | )]
| | size = (18, 2)
| |
num_edges_per_mesh | size = (N) | tensor([3, 5, 10])
| | size = (3)
| |
edges_packed_to_mesh_idx | size = (sum(E_n)) | tensor([
| | 0, 0, 0,
| | . . .
| | 2, 2, 2
| | ])
| | size = (18)
| |
faces_packed_to_edges_packed | size = (sum(F_n), 3) | tensor([
| | [2, 1, 0],
| | [5, 4, 3],
| | . . .
| | [12, 14, 16],
| | ])
| | size = (10, 3)
| |
----------------------------------------------------------------------------
# SPHINX IGNORE
"""
_INTERNAL_TENSORS = [
"_verts_packed",
"_verts_packed_to_mesh_idx",
"_mesh_to_verts_packed_first_idx",
"_verts_padded",
"_num_verts_per_mesh",
"_faces_packed",
"_faces_packed_to_mesh_idx",
"_mesh_to_faces_packed_first_idx",
"_faces_padded",
"_faces_areas_packed",
"_verts_normals_packed",
"_faces_normals_packed",
"_num_faces_per_mesh",
"_edges_packed",
"_edges_packed_to_mesh_idx",
"_faces_packed_to_edges_packed",
"_num_edges_per_mesh",
"_verts_padded_to_packed_idx",
"_laplacian_packed",
"valid",
"equisized",
]
def __init__(self, verts=None, faces=None, textures=None):
"""
Args:
verts:
Can be either
- List where each element is a tensor of shape (num_verts, 3)
containing the (x, y, z) coordinates of each vertex.
- Padded float tensor with shape (num_meshes, max_num_verts, 3).
Meshes should be padded with fill value of 0 so they all have
the same number of vertices.
faces:
Can be either
- List where each element is a tensor of shape (num_faces, 3)
containing the indices of the 3 vertices in the corresponding
mesh in verts which form the triangular face.
- Padded long tensor of shape (num_meshes, max_num_faces, 3).
Meshes should be padded with fill value of -1 so they have
the same number of faces.
textures: Optional instance of the Textures class with mesh
texture properties.
Refer to comments above for descriptions of List and Padded representations.
"""
self.device = None
if textures is not None and not isinstance(textures, Textures):
msg = "Expected textures to be of type Textures; got %r"
raise ValueError(msg % type(textures))
self.textures = textures
# Indicates whether the meshes in the list/batch have the same number
# of faces and vertices.
self.equisized = False
# Boolean indicator for each mesh in the batch
# True if mesh has non zero number of verts and face, False otherwise.
self.valid = None
self._N = 0 # batch size (number of meshes)
self._V = 0 # (max) number of vertices per mesh
self._F = 0 # (max) number of faces per mesh
# List of Tensors of verts and faces.
self._verts_list = None
self._faces_list = None
# Packed representation for verts.
self._verts_packed = None # (sum(V_n), 3)
self._verts_packed_to_mesh_idx = None # sum(V_n)
# Index to convert verts from flattened padded to packed
self._verts_padded_to_packed_idx = None # N * max_V
# Index of each mesh's first vert in the packed verts.
# Assumes packing is sequential.
self._mesh_to_verts_packed_first_idx = None # N
# Packed representation for faces.
self._faces_packed = None # (sum(F_n), 3)
self._faces_packed_to_mesh_idx = None # sum(F_n)
# Index of each mesh's first face in packed faces.
# Assumes packing is sequential.
self._mesh_to_faces_packed_first_idx = None # N
# Packed representation of edges sorted by index of the first vertex
# in the edge. Edges can be shared between faces in a mesh.
self._edges_packed = None # (sum(E_n), 2)
# Map from packed edges to corresponding mesh index.
self._edges_packed_to_mesh_idx = None # sum(E_n)
self._num_edges_per_mesh = None # N
# Map from packed faces to packed edges. This represents the index of
# the edge opposite the vertex for each vertex in the face. E.g.
#
# v0
# /\
# / \
# e1 / \ e2
# / \
# /________\
# v2 e0 v1
#
# Face (v0, v1, v2) => Edges (e0, e1, e2)
self._faces_packed_to_edges_packed = None # (sum(F_n), 3)
# Padded representation of verts.
self._verts_padded = None # (N, max(V_n), 3)
self._num_verts_per_mesh = None # N
# Padded representation of faces.
self._faces_padded = None # (N, max(F_n), 3)
self._num_faces_per_mesh = None # N
# Face areas
self._faces_areas_packed = None
# Normals
self._verts_normals_packed = None
self._faces_normals_packed = None
# Packed representation of Laplacian Matrix
self._laplacian_packed = None
# Identify type of verts and faces.
if isinstance(verts, list) and isinstance(faces, list):
self._verts_list = verts
self._faces_list = [
f[f.gt(-1).all(1)].to(torch.int64) if len(f) > 0 else f
for f in faces
]
self._N = len(self._verts_list)
self.device = torch.device("cpu")
self.valid = torch.zeros(
(self._N,), dtype=torch.bool, device=self.device
)
if self._N > 0:
self.device = self._verts_list[0].device
num_verts_per_mesh = torch.tensor(
[len(v) for v in self._verts_list], device=self.device
)
self._V = num_verts_per_mesh.max()
num_faces_per_mesh = torch.tensor(
[len(f) for f in self._faces_list], device=self.device
)
self._F = num_faces_per_mesh.max()
self.valid = torch.tensor(
[
len(v) > 0 and len(f) > 0
for (v, f) in zip(self._verts_list, self._faces_list)
],
dtype=torch.bool,
device=self.device,
)
if (len(num_verts_per_mesh.unique()) == 1) and (
len(num_faces_per_mesh.unique()) == 1
):
self.equisized = True
elif torch.is_tensor(verts) and torch.is_tensor(faces):
if verts.size(2) != 3 and faces.size(2) != 3:
raise ValueError(
"Verts and Faces tensors have incorrect dimensions."
)
self._verts_padded = verts
self._faces_padded = faces.to(torch.int64)
self._N = self._verts_padded.shape[0]
self._V = self._verts_padded.shape[1]
self.device = self._verts_padded.device
self.valid = torch.zeros(
(self._N,), dtype=torch.bool, device=self.device
)
if self._N > 0:
# Check that padded faces - which have value -1 - are at the
# end of the tensors
faces_not_padded = self._faces_padded.gt(-1).all(2)
num_faces = faces_not_padded.sum(1)
if (faces_not_padded[:, :-1] < faces_not_padded[:, 1:]).any():
raise ValueError("Padding of faces must be at the end")
# NOTE that we don't check for the ordering of padded verts
# as long as the faces index correspond to the right vertices.
self.valid = num_faces > 0
self._F = num_faces.max()
if len(num_faces.unique()) == 1:
self.equisized = True
else:
raise ValueError(
"Verts and Faces must be either a list or a tensor with \
shape (batch_size, N, 3) where N is either the maximum \
number of verts or faces respectively."
)
def __len__(self):
return self._N
def __getitem__(self, index):
"""
Args:
index: Specifying the index of the mesh to retrieve.
Can be an int, slice, list of ints or a boolean tensor.
Returns:
Meshes object with selected meshes. The mesh tensors are not cloned.
"""
if isinstance(index, (int, slice)):
verts = self.verts_list()[index]
faces = self.faces_list()[index]
elif isinstance(index, list):
verts = [self.verts_list()[i] for i in index]
faces = [self.faces_list()[i] for i in index]
elif isinstance(index, torch.Tensor):
if index.dim() != 1 or index.dtype.is_floating_point:
raise IndexError(index)
# NOTE consider converting index to cpu for efficiency
if index.dtype == torch.bool:
# advanced indexing on a single dimension
index = index.nonzero()
index = index.squeeze(1) if index.numel() > 0 else index
index = index.tolist()
verts = [self.verts_list()[i] for i in index]
faces = [self.faces_list()[i] for i in index]
else:
raise IndexError(index)
if torch.is_tensor(verts) and torch.is_tensor(faces):
return Meshes(verts=[verts], faces=[faces])
elif isinstance(verts, list) and isinstance(faces, list):
return Meshes(verts=verts, faces=faces)
else:
raise ValueError("(verts, faces) not defined correctly")
def isempty(self) -> bool:
"""
Checks whether any mesh is valid.
Returns:
bool indicating whether there is any data.
"""
return self._N == 0 or self.valid.eq(False).all()
def verts_list(self):
"""
Get the list representation of the vertices.
Returns:
list of tensors of vertices of shape (V_n, 3).
"""
if self._verts_list is None:
assert (
self._verts_padded is not None
), "verts_padded is required to compute verts_list."
self._verts_list = [
v[0] for v in self._verts_padded.split([1] * self._N, 0)
]
return self._verts_list
def faces_list(self):
"""
Get the list representation of the faces.
Returns:
list of tensors of faces of shape (F_n, 3).
"""
if self._faces_list is None:
assert (
self._faces_padded is not None
), "faces_padded is required to compute faces_list."
self._faces_list = []
for i in range(self._N):
valid = self._faces_padded[i].gt(-1).all(1)
self._faces_list.append(self._faces_padded[i, valid, :])
return self._faces_list
def verts_packed(self):
"""
Get the packed representation of the vertices.
Returns:
tensor of vertices of shape (sum(V_n), 3).
"""
self._compute_packed()
return self._verts_packed
def verts_packed_to_mesh_idx(self):
"""
Return a 1D tensor with the same first dimension as verts_packed.
verts_packed_to_mesh_idx[i] gives the index of the mesh which contains
verts_packed[i].
Returns:
1D tensor of indices.
"""
self._compute_packed()
return self._verts_packed_to_mesh_idx
def mesh_to_verts_packed_first_idx(self):
"""
Return a 1D tensor x with length equal to the number of meshes such that
the first vertex of the ith mesh is verts_packed[x[i]].
Returns:
1D tensor of indices of first items.
"""
self._compute_packed()
return self._mesh_to_verts_packed_first_idx
def num_verts_per_mesh(self):
"""
Return a 1D tensor x with length equal to the number of meshes giving
the number of vertices in each mesh.
Returns:
1D tensor of sizes.
"""
self._compute_packed()
return self._num_verts_per_mesh
def faces_packed(self):
"""
Get the packed representation of the faces.
Faces are given by the indices of the three vertices in verts_packed.
Returns:
tensor of faces of shape (sum(F_n), 3).
"""
self._compute_packed()
return self._faces_packed
def faces_packed_to_mesh_idx(self):
"""
Return a 1D tensor with the same first dimension as faces_packed.
faces_packed_to_mesh_idx[i] gives the index of the mesh which contains
faces_packed[i].
Returns:
1D tensor of indices.
"""
self._compute_packed()
return self._faces_packed_to_mesh_idx
def mesh_to_faces_packed_first_idx(self):
"""
Return a 1D tensor x with length equal to the number of meshes such that
the first face of the ith mesh is faces_packed[x[i]].
Returns:
1D tensor of indices of first items.
"""
self._compute_packed()
return self._mesh_to_faces_packed_first_idx
def verts_padded(self):
"""
Get the padded representation of the vertices.
Returns:
tensor of vertices of shape (N, max(V_n), 3).
"""
self._compute_padded()
return self._verts_padded
def faces_padded(self):
"""
Get the padded representation of the faces.
Returns:
tensor of faces of shape (N, max(F_n), 3).
"""
self._compute_padded()
return self._faces_padded
def num_faces_per_mesh(self):
"""
Return a 1D tensor x with length equal to the number of meshes giving
the number of faces in each mesh.
Returns:
1D tensor of sizes.
"""
self._compute_packed()
return self._num_faces_per_mesh
def edges_packed(self):
"""
Get the packed representation of the edges.
Returns:
tensor of edges of shape (sum(E_n), 2).
"""
self._compute_edges_packed()
return self._edges_packed
def edges_packed_to_mesh_idx(self):
"""
Return a 1D tensor with the same first dimension as edges_packed.
edges_packed_to_mesh_idx[i] gives the index of the mesh which contains
edges_packed[i].
Returns:
1D tensor of indices.
"""
self._compute_edges_packed()
return self._edges_packed_to_mesh_idx
def faces_packed_to_edges_packed(self):
"""
Get the packed representation of the faces in terms of edges.
Faces are given by the indices of the three edges in
the packed representation of the edges.
Returns:
tensor of faces of shape (sum(F_n), 3).
"""
self._compute_edges_packed()
return self._faces_packed_to_edges_packed
def num_edges_per_mesh(self):
"""
Return a 1D tensor x with length equal to the number of meshes giving
the number of edges in each mesh.
Returns:
1D tensor of sizes.
"""
self._compute_edges_packed()
return self._num_edges_per_mesh
def verts_padded_to_packed_idx(self):
"""
Return a 1D tensor x with length equal to the total number of vertices
such that verts_packed()[i] is element x[i] of the flattened padded
representation.
The packed representation can be calculated as follows.
.. code-block:: python
p = verts_padded().reshape(-1, 3)
verts_packed = p[x]
Returns:
1D tensor of indices.
"""
self._compute_packed()
if self._verts_padded_to_packed_idx is not None:
return self._verts_padded_to_packed_idx
self._verts_padded_to_packed_idx = torch.cat(
[
torch.arange(v, dtype=torch.int64, device=self.device)
+ i * self._V
for (i, v) in enumerate(self._num_verts_per_mesh)
],
dim=0,
)
return self._verts_padded_to_packed_idx
def verts_normals_packed(self):
"""
Get the packed representation of the vertex normals.
Returns:
tensor of normals of shape (sum(V_n), 3).
"""
self._compute_vertex_normals()
return self._verts_normals_packed
def verts_normals_list(self):
"""
Get the list representation of the vertex normals.
Returns:
list of tensors of normals of shape (V_n, 3).
"""
if self.isempty():
return [
torch.empty((0, 3), dtype=torch.float32, device=self.device)
] * self._N
verts_normals_packed = self.verts_normals_packed()
split_size = self.num_verts_per_mesh().tolist()
return struct_utils.packed_to_list(verts_normals_packed, split_size)
def verts_normals_padded(self):
"""
Get the padded representation of the vertex normals.
Returns:
tensor of normals of shape (N, max(V_n), 3).
"""
if self.isempty():
return torch.zeros(
(self._N, 0, 3), dtype=torch.float32, device=self.device
)
verts_normals_list = self.verts_normals_list()
return struct_utils.list_to_padded(
verts_normals_list,
(self._V, 3),
pad_value=0.0,
equisized=self.equisized,
)
def faces_normals_packed(self):
"""
Get the packed representation of the face normals.
Returns:
tensor of normals of shape (sum(F_n), 3).
"""
self._compute_face_areas_normals()
return self._faces_normals_packed
def faces_normals_list(self):
"""
Get the list representation of the face normals.
Returns:
list of tensors of normals of shape (F_n, 3).
"""
if self.isempty():
return [
torch.empty((0, 3), dtype=torch.float32, device=self.device)
] * self._N
faces_normals_packed = self.faces_normals_packed()
split_size = self.num_faces_per_mesh().tolist()
return struct_utils.packed_to_list(faces_normals_packed, split_size)
def faces_normals_padded(self):
"""
Get the padded representation of the face normals.
Returns:
tensor of normals of shape (N, max(F_n), 3).
"""
if self.isempty():
return torch.zeros(
(self._N, 0, 3), dtype=torch.float32, device=self.device
)
faces_normals_list = self.faces_normals_list()
return struct_utils.list_to_padded(
faces_normals_list,
(self._F, 3),
pad_value=0.0,
equisized=self.equisized,
)
def faces_areas_packed(self):
"""
Get the packed representation of the face areas.
Returns:
tensor of areas of shape (sum(F_n),).
"""
self._compute_face_areas_normals()
return self._faces_areas_packed
def laplacian_packed(self):
self._compute_laplacian_packed()
return self._laplacian_packed
def _compute_face_areas_normals(self, refresh: bool = False):
"""
Compute the area and normal of each face in faces_packed.
The convention of a normal for a face consisting of verts [v0, v1, v2]
is normal = (v1 - v0) x (v2 - v0)
Args:
refresh: Set to True to force recomputation of face areas.
Default: False.
"""
if not (
refresh
or any(
v is None
for v in [self._faces_areas_packed, self._faces_normals_packed]
)
):
return
faces_packed = self.faces_packed()
verts_packed = self.verts_packed()
if verts_packed.is_cuda and faces_packed.is_cuda:
face_areas, face_normals = _C.face_areas_normals(
verts_packed, faces_packed
)
else:
vertices_faces = verts_packed[faces_packed] # (F, 3, 3)
# vector pointing from v0 to v1
v01 = vertices_faces[:, 1] - vertices_faces[:, 0]
# vector pointing from v0 to v2
v02 = vertices_faces[:, 2] - vertices_faces[:, 0]
normals = torch.cross(v01, v02, dim=1) # (F, 3)
face_areas = normals.norm(dim=-1) / 2
face_normals = torch.nn.functional.normalize(
normals, p=2, dim=1, eps=1e-6
)
self._faces_areas_packed = face_areas
self._faces_normals_packed = face_normals
def _compute_vertex_normals(self, refresh: bool = False):
"""Computes the packed version of vertex normals from the packed verts
and faces. This assumes verts are shared between faces. The normal for
a vertex is computed as the sum of the normals of all the faces it is
part of weighed by the face areas.
Args:
refresh: Set to True to force recomputation of vertex normals.
Default: False.
"""
if not (
refresh or any(v is None for v in [self._verts_normals_packed])
):
return
if self.isempty():
self._verts_normals_packed = torch.zeros(
(self._N, 3), dtype=torch.int64, device=self.device
)
else:
faces_packed = self.faces_packed()
verts_packed = self.verts_packed()
verts_normals = torch.zeros_like(verts_packed)
vertices_faces = verts_packed[faces_packed]
# NOTE: this is already applying the area weighting as the magnitude
# of the cross product is 2 x area of the triangle.
verts_normals = verts_normals.index_add(
0,
faces_packed[:, 1],
torch.cross(
vertices_faces[:, 2] - vertices_faces[:, 1],
vertices_faces[:, 0] - vertices_faces[:, 1],
dim=1,
),
)
verts_normals = verts_normals.index_add(
0,
faces_packed[:, 2],
torch.cross(
vertices_faces[:, 0] - vertices_faces[:, 2],
vertices_faces[:, 1] - vertices_faces[:, 2],
dim=1,
),
)
verts_normals = verts_normals.index_add(
0,
faces_packed[:, 0],
torch.cross(
vertices_faces[:, 1] - vertices_faces[:, 0],
vertices_faces[:, 2] - vertices_faces[:, 0],
dim=1,
),
)
self._verts_normals_packed = torch.nn.functional.normalize(
verts_normals, eps=1e-6, dim=1
)
def _compute_padded(self, refresh: bool = False):
"""
Computes the padded version of meshes from verts_list and faces_list.
"""
if not (
refresh
or any(v is None for v in [self._verts_padded, self._faces_padded])
):
return
verts_list = self._verts_list
faces_list = self._faces_list
assert (
faces_list is not None and verts_list is not None
), "faces_list and verts_list arguments are required"
if self.isempty():
self._faces_padded = torch.zeros(
(self._N, 0, 3), dtype=torch.int64, device=self.device
)
self._verts_padded = torch.zeros(
(self._N, 0, 3), dtype=torch.float32, device=self.device
)
else:
self._faces_padded = struct_utils.list_to_padded(
faces_list,
(self._F, 3),
pad_value=-1.0,
equisized=self.equisized,
)
self._verts_padded = struct_utils.list_to_padded(
verts_list,
(self._V, 3),
pad_value=0.0,
equisized=self.equisized,
)
# TODO(nikhilar) Improve performance of _compute_packed.
def _compute_packed(self, refresh: bool = False):
"""
Computes the packed version of the meshes from verts_list and faces_list
and sets the values of auxillary tensors.
Args:
refresh: Set to True to force recomputation of packed representations.
Default: False.
"""
if not (
refresh
or any(
v is None
for v in [
self._verts_packed,
self._verts_packed_to_mesh_idx,
self._mesh_to_verts_packed_first_idx,
self._num_verts_per_mesh,
self._faces_packed,
self._faces_packed_to_mesh_idx,
self._mesh_to_faces_packed_first_idx,
self._num_faces_per_mesh,
]
)
):
return
# Packed can be calculated from padded or list, so can call the
# accessor function for verts_list and faces_list.
verts_list = self.verts_list()
faces_list = self.faces_list()
if self.isempty():
self._verts_packed = torch.zeros(
(0, 3), dtype=torch.float32, device=self.device
)
self._verts_packed_to_mesh_idx = torch.zeros(
(0,), dtype=torch.int64, device=self.device
)
self._mesh_to_verts_packed_first_idx = torch.zeros(
(0,), dtype=torch.int64, device=self.device
)
self._num_verts_per_mesh = torch.zeros(
(0,), dtype=torch.int64, device=self.device
)
self._faces_packed = -torch.ones(
(0, 3), dtype=torch.int64, device=self.device
)
self._faces_packed_to_mesh_idx = torch.zeros(
(0,), dtype=torch.int64, device=self.device
)
self._mesh_to_faces_packed_first_idx = torch.zeros(
(0,), dtype=torch.int64, device=self.device
)
self._num_faces_per_mesh = torch.zeros(
(0,), dtype=torch.int64, device=self.device
)
return
verts_list_to_packed = struct_utils.list_to_packed(verts_list)
self._verts_packed = verts_list_to_packed[0]
self._num_verts_per_mesh = verts_list_to_packed[1]
self._mesh_to_verts_packed_first_idx = verts_list_to_packed[2]
self._verts_packed_to_mesh_idx = verts_list_to_packed[3]
faces_list_to_packed = struct_utils.list_to_packed(faces_list)
faces_packed = faces_list_to_packed[0]
self._num_faces_per_mesh = faces_list_to_packed[1]
self._mesh_to_faces_packed_first_idx = faces_list_to_packed[2]
self._faces_packed_to_mesh_idx = faces_list_to_packed[3]
faces_packed_offset = self._mesh_to_verts_packed_first_idx[
self._faces_packed_to_mesh_idx
]
self._faces_packed = faces_packed + faces_packed_offset.view(-1, 1)
def _compute_edges_packed(self, refresh: bool = False):
"""
Computes edges in packed form from the packed version of faces and verts.
"""
if not (
refresh
or any(
v is None
for v in [
self._edges_packed,
self._faces_packed_to_mesh_idx,
self._edges_packed_to_mesh_idx,
self._num_edges_per_mesh,
]
)
):
return
if self.isempty():
self._edges_packed = -torch.ones(
(0, 2), dtype=torch.int64, device=self.device
)
self._edges_packed_to_mesh_idx = torch.zeros(
(0,), dtype=torch.int64, device=self.device
)
return
faces = self.faces_packed()
F = faces.shape[0]
v0, v1, v2 = faces.chunk(3, dim=1)
e01 = torch.cat([v0, v1], dim=1) # (sum(F_n), 2)
e12 = torch.cat([v1, v2], dim=1) # (sum(F_n), 2)
e20 = torch.cat([v2, v0], dim=1) # (sum(F_n), 2)
# All edges including duplicates.
edges = torch.cat([e12, e20, e01], dim=0) # (sum(F_n)*3, 2)
edge_to_mesh = torch.cat(
[
self._faces_packed_to_mesh_idx,
self._faces_packed_to_mesh_idx,
self._faces_packed_to_mesh_idx,
],
dim=0,
) # sum(F_n)*3
# Sort the edges in increasing vertex order to remove duplicates as
# the same edge may appear in different orientations in different faces.
# i.e. rows in edges after sorting will be of the form (v0, v1) where v1 > v0.
# This sorting does not change the order in dim=0.
edges, _ = edges.sort(dim=1)
# Remove duplicate edges: convert each edge (v0, v1) into an
# integer hash = V * v0 + v1; this allows us to use the scalar version of
# unique which is much faster than edges.unique(dim=1) which is very slow.
# After finding the unique elements reconstruct the vertex indicies as:
# (v0, v1) = (hash / V, hash % V)
# The inverse maps from unique_edges back to edges:
# unique_edges[inverse_idxs] == edges
# i.e. inverse_idxs[i] == j means that edges[i] == unique_edges[j]
V = self._verts_packed.shape[0]
edges_hash = V * edges[:, 0] + edges[:, 1]
u, inverse_idxs = torch.unique(edges_hash, return_inverse=True)
# Find indices of unique elements.
# TODO (nikhilar) remove following 4 lines when torch.unique has support
# for returning unique indices
sorted_hash, sort_idx = torch.sort(edges_hash, dim=0)
unique_mask = torch.ones(
edges_hash.shape[0], dtype=torch.bool, device=self.device
)
unique_mask[1:] = sorted_hash[1:] != sorted_hash[:-1]
unique_idx = sort_idx[unique_mask]
self._edges_packed = torch.stack([u / V, u % V], dim=1)
self._edges_packed_to_mesh_idx = edge_to_mesh[unique_idx]
face_to_edge = torch.arange(3 * F).view(3, F).t()
face_to_edge = inverse_idxs[face_to_edge]
self._faces_packed_to_edges_packed = face_to_edge
num_edges_per_mesh = torch.zeros(
self._N, dtype=torch.int32, device=self.device
)
ones = torch.ones(1, dtype=torch.int32, device=self.device).expand(
self._edges_packed_to_mesh_idx.shape
)
self._num_edges_per_mesh = num_edges_per_mesh.scatter_add(
0, self._edges_packed_to_mesh_idx, ones
)
def _compute_laplacian_packed(self, refresh: bool = False):
"""
Computes the laplacian in packed form.
The definition of the laplacian is
L[i, j] = -1 , if i == j
L[i, j] = 1 / deg(i) , if (i, j) is an edge
L[i, j] = 0 , otherwise
where deg(i) is the degree of the i-th vertex in the graph
Returns:
Sparse FloatTensor of shape (V, V) where V = sum(V_n)
"""
if not (refresh or self._laplacian_packed is None):
return
if self.isempty():
self._laplacian_packed = torch.zeros(
(0, 0), dtype=torch.float32, device=self.device
).to_sparse()
return
verts_packed = self.verts_packed() # (sum(V_n), 3)
edges_packed = self.edges_packed() # (sum(E_n), 3)
V = verts_packed.shape[0] # sum(V_n)
e0, e1 = edges_packed.unbind(1)
idx01 = torch.stack([e0, e1], dim=1) # (sum(E_n), 2)
idx10 = torch.stack([e1, e0], dim=1) # (sum(E_n), 2)
idx = torch.cat([idx01, idx10], dim=0).t() # (2, 2*sum(E_n))
# First, we construct the adjacency matrix,
# i.e. A[i, j] = 1 if (i,j) is an edge, or
# A[e0, e1] = 1 & A[e1, e0] = 1
ones = torch.ones(idx.shape[1], dtype=torch.float32, device=self.device)
A = torch.sparse.FloatTensor(idx, ones, (V, V))
# the sum of i-th row of A gives the degree of the i-th vertex
deg = torch.sparse.sum(A, dim=1).to_dense()
# We construct the Laplacian matrix by adding the non diagonal values
# i.e. L[i, j] = 1 ./ deg(i) if (i, j) is an edge
deg0 = deg[e0]
deg0 = torch.where(deg0 > 0.0, 1.0 / deg0, deg0)
deg1 = deg[e1]
deg1 = torch.where(deg1 > 0.0, 1.0 / deg1, deg1)
val = torch.cat([deg0, deg1])
L = torch.sparse.FloatTensor(idx, val, (V, V))
# Then we add the diagonal values L[i, i] = -1.
idx = torch.arange(V, device=self.device)
idx = torch.stack([idx, idx], dim=0)
ones = torch.ones(idx.shape[1], dtype=torch.float32, device=self.device)
L -= torch.sparse.FloatTensor(idx, ones, (V, V))
self._laplacian_packed = L
def clone(self):
"""
Deep copy of Meshes object. All internal tensors are cloned individually.
Returns:
new Meshes object.
"""
verts_list = self.verts_list()
faces_list = self.faces_list()
new_verts_list = [v.clone() for v in verts_list]
new_faces_list = [f.clone() for f in faces_list]
other = Meshes(verts=new_verts_list, faces=new_faces_list)
for k in self._INTERNAL_TENSORS:
v = getattr(self, k)
if torch.is_tensor(v):
setattr(other, k, v.clone())
# Textures is not a tensor but has a clone method
if self.textures is not None:
other.textures = self.textures.clone()
return other
def to(self, device, copy: bool = False):
"""
Match functionality of torch.Tensor.to()
If copy = True or the self Tensor is on a different device, the
returned tensor is a copy of self with the desired torch.device.
If copy = False and the self Tensor already has the correct torch.device,
then self is returned.
Args:
device: Device id for the new tensor.
copy: Boolean indicator whether or not to clone self. Default False.
Returns:
Meshes object.
"""
if not copy and self.device == device:
return self
other = self.clone()
if self.device != device:
other.device = device
if other._N > 0:
other._verts_list = [v.to(device) for v in other._verts_list]
other._faces_list = [f.to(device) for f in other._faces_list]
for k in self._INTERNAL_TENSORS:
v = getattr(self, k)
if torch.is_tensor(v):
setattr(other, k, v.to(device))
if self.textures is not None:
other.textures = self.textures.to(device)
return other
def cpu(self):
return self.to(torch.device("cpu"))
def cuda(self):
return self.to(torch.device("cuda"))
def get_mesh_verts_faces(self, index: int):
"""
Get tensors for a single mesh from the list representation.
Args:
index: Integer in the range [0, N).
Returns:
verts: Tensor of shape (V, 3).
faces: LongTensor of shape (F, 3).
"""
if not isinstance(index, int):
raise ValueError("Mesh index must be an integer.")
if index < 0 or index > self._N:
raise ValueError(
"Mesh index must be in the range [0, N) where \
N is the number of meshes in the batch."
)
verts = self.verts_list()
faces = self.faces_list()
return verts[index], faces[index]
# TODO(nikhilar) Move function to a utils file.
def split(self, split_sizes: list):
"""
Splits Meshes object of size N into a list of Meshes objects of
size len(split_sizes), where the i-th Meshes object is of size split_sizes[i].
Similar to torch.split().
Args:
split_sizes: List of integer sizes of Meshes objects to be returned.
Returns:
list[Meshes].
"""
if not all(isinstance(x, int) for x in split_sizes):
raise ValueError("Value of split_sizes must be a list of integers.")
meshlist = []
curi = 0
for i in split_sizes:
meshlist.append(self[curi : curi + i])
curi += i
return meshlist
def offset_verts_(self, vert_offsets_packed):
"""
Add an offset to the vertices of this Meshes. In place operation.
Args:
vert_offsets_packed: A Tensor of the same shape as self.verts_packed
giving offsets to be added to all vertices.
Returns:
self.
"""
verts_packed = self.verts_packed()
if vert_offsets_packed.shape != verts_packed.shape:
raise ValueError("Verts offsets must have dimension (all_v, 2).")
# update verts packed
self._verts_packed = verts_packed + vert_offsets_packed
new_verts_list = list(
self._verts_packed.split(self.num_verts_per_mesh().tolist(), 0)
)
# update verts list
# Note that since _compute_packed() has been executed, verts_list
# cannot be None even if not provided during construction.
self._verts_list = new_verts_list
# update verts padded
if self._verts_padded is not None:
for i, verts in enumerate(new_verts_list):
if len(verts) > 0:
self._verts_padded[i, : verts.shape[0], :] = verts
# update face areas and normals and vertex normals
# only if the original attributes are computed
if any(
v is not None
for v in [self._faces_areas_packed, self._faces_normals_packed]
):
self._compute_face_areas_normals(refresh=True)
if self._verts_normals_packed is not None:
self._compute_vertex_normals(refresh=True)
return self
# TODO(nikhilar) Move out of place operator to a utils file.
def offset_verts(self, vert_offsets_packed):
"""
Out of place offset_verts.
Args:
vert_offsets_packed: A Tensor of the same shape as self.verts_packed
giving offsets to be added to all vertices.
Returns:
new Meshes object.
"""
new_mesh = self.clone()
return new_mesh.offset_verts_(vert_offsets_packed)
def scale_verts_(self, scale):
"""
Multiply the vertices of this Meshes object by a scalar value.
In place operation.
Args:
scale: A scalar, or a Tensor of shape (N,).
Returns:
self.
"""
if not torch.is_tensor(scale):
scale = torch.full((len(self),), scale, device=self.device)
new_verts_list = []
verts_list = self.verts_list()
for i, old_verts in enumerate(verts_list):
new_verts_list.append(scale[i] * old_verts)
# update list
self._verts_list = new_verts_list
# update packed
if self._verts_packed is not None:
self._verts_packed = torch.cat(new_verts_list, dim=0)
# update padded
if self._verts_padded is not None:
for i, verts in enumerate(self._verts_list):
if len(verts) > 0:
self._verts_padded[i, : verts.shape[0], :] = verts
# update face areas and normals and vertex normals
# only if the original attributes are computed
if any(
v is not None
for v in [self._faces_areas_packed, self._faces_normals_packed]
):
self._compute_face_areas_normals(refresh=True)
if self._verts_normals_packed is not None:
self._compute_vertex_normals(refresh=True)
return self
def scale_verts(self, scale):
"""
Out of place scale_verts.
Args:
scale: A scalar, or a Tensor of shape (N,).
Returns:
new Meshes object.
"""
new_mesh = self.clone()
return new_mesh.scale_verts_(scale)
# TODO(nikhilar) Move function to utils file.
def get_bounding_boxes(self):
"""
Compute an axis-aligned bounding box for each mesh in this Meshes object.
Returns:
bboxes: Tensor of shape (N, 3, 2) where bbox[i, j] gives the
min and max values of mesh i along the jth coordinate axis.
"""
all_mins, all_maxes = [], []
for verts in self.verts_list():
cur_mins = verts.min(dim=0)[0] # (3,)
cur_maxes = verts.max(dim=0)[0] # (3,)
all_mins.append(cur_mins)
all_maxes.append(cur_maxes)
all_mins = torch.stack(all_mins, dim=0) # (N, 3)
all_maxes = torch.stack(all_maxes, dim=0) # (N, 3)
bboxes = torch.stack([all_mins, all_maxes], dim=2)
return bboxes
def extend(self, N: int):
"""
Create new Meshes class which contains each input mesh N times
Args:
N: number of new copies of each mesh.
Returns:
new Meshes object.
"""
if not isinstance(N, int):
raise ValueError("N must be an integer.")
if N <= 0:
raise ValueError("N must be > 0.")
new_verts_list, new_faces_list = [], []
for verts, faces in zip(self.verts_list(), self.faces_list()):
new_verts_list.extend(verts.clone() for _ in range(N))
new_faces_list.extend(faces.clone() for _ in range(N))
tex = None
if self.textures is not None:
tex = self.textures.extend(N)
return Meshes(verts=new_verts_list, faces=new_faces_list, textures=tex)
pytorch3d/structures/textures.py 0 → 100644
View file @ dbf06b50
#!/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)
pytorch3d/structures/utils.py 0 → 100644
View file @ dbf06b50
#!/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)
pytorch3d/transforms/__init__.py 0 → 100644
View file @ dbf06b50
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
from .rotation_conversions import (
euler_angles_to_matrix,
matrix_to_euler_angles,
matrix_to_quaternion,
quaternion_apply,
quaternion_invert,
quaternion_multiply,
quaternion_raw_multiply,
quaternion_to_matrix,
random_quaternions,
random_rotation,
random_rotations,
standardize_quaternion,
)
from .so3 import (
so3_exponential_map,
so3_log_map,
so3_relative_angle,
so3_rotation_angle,
)
from .transform3d import Rotate, RotateAxisAngle, Scale, Transform3d, Translate
__all__ = [k for k in globals().keys() if not k.startswith("_")]
pytorch3d/transforms/rotation_conversions.py 0 → 100644
View file @ dbf06b50
#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
import functools
import torch
def quaternion_to_matrix(quaternions):
"""
Convert rotations given as quaternions to rotation matrices.
Args:
quaternions: quaternions with real part first,
as tensor of shape (..., 4).
Returns:
Rotation matrices as tensor of shape (..., 3, 3).
"""
r, i, j, k = torch.unbind(quaternions, -1)
two_s = 2.0 / (quaternions * quaternions).sum(-1)
o = torch.stack(
(
1 - two_s * (j * j + k * k),
two_s * (i * j - k * r),
two_s * (i * k + j * r),
two_s * (i * j + k * r),
1 - two_s * (i * i + k * k),
two_s * (j * k - i * r),
two_s * (i * k - j * r),
two_s * (j * k + i * r),
1 - two_s * (i * i + j * j),
),
-1,
)
return o.reshape(quaternions.shape[:-1] + (3, 3))
def _copysign(a, b):
"""
Return a tensor where each element has the absolute value taken from the,
corresponding element of a, with sign taken from the corresponding
element of b. This is like the standard copysign floating-point operation,
but is not careful about negative 0 and NaN.
Args:
a: source tensor.
b: tensor whose signs will be used, of the same shape as a.
Returns:
Tensor of the same shape as a with the signs of b.
"""
signs_differ = (a < 0) != (b < 0)
return torch.where(signs_differ, -a, a)
def matrix_to_quaternion(matrix):
"""
Convert rotations given as rotation matrices to quaternions.
Args:
matrix: Rotation matrices as tensor of shape (..., 3, 3).
Returns:
quaternions with real part first, as tensor of shape (..., 4).
"""
if matrix.size(-1) != 3 or matrix.size(-2) != 3:
raise ValueError(f"Invalid rotation matrix shape f{matrix.shape}.")
zero = matrix.new_zeros((1,))
m00 = matrix[..., 0, 0]
m11 = matrix[..., 1, 1]
m22 = matrix[..., 2, 2]
o0 = 0.5 * torch.sqrt(torch.max(zero, 1 + m00 + m11 + m22))
x = 0.5 * torch.sqrt(torch.max(zero, 1 + m00 - m11 - m22))
y = 0.5 * torch.sqrt(torch.max(zero, 1 - m00 + m11 - m22))
z = 0.5 * torch.sqrt(torch.max(zero, 1 - m00 - m11 + m22))
o1 = _copysign(x, matrix[..., 2, 1] - matrix[..., 1, 2])
o2 = _copysign(y, matrix[..., 0, 2] - matrix[..., 2, 0])
o3 = _copysign(z, matrix[..., 1, 0] - matrix[..., 0, 1])
return torch.stack((o0, o1, o2, o3), -1)
def _primary_matrix(axis: str, angle):
"""
Return the rotation matrices for one of the rotations about an axis
of which Euler angles describe, for each value of the angle given.
Args:
axis: Axis label "X" or "Y or "Z".
angle: any shape tensor of Euler angles in radians
Returns:
Rotation matrices as tensor of shape (..., 3, 3).
"""
cos = torch.cos(angle)
sin = torch.sin(angle)
one = torch.ones_like(angle)
zero = torch.zeros_like(angle)
if axis == "X":
o = (one, zero, zero, zero, cos, -sin, zero, sin, cos)
if axis == "Y":
o = (cos, zero, sin, zero, one, zero, -sin, zero, cos)
if axis == "Z":
o = (cos, -sin, zero, sin, cos, zero, zero, zero, one)
return torch.stack(o, -1).reshape(angle.shape + (3, 3))
def euler_angles_to_matrix(euler_angles, convention: str):
"""
Convert rotations given as Euler angles in radians to rotation matrices.
Args:
euler_angles: Euler angles in radians as tensor of shape (..., 3).
convention: Convention string of three uppercase letters from
{"X", "Y", and "Z"}.
Returns:
Rotation matrices as tensor of shape (..., 3, 3).
"""
if euler_angles.dim() == 0 or euler_angles.shape[-1] != 3:
raise ValueError("Invalid input euler angles.")
if len(convention) != 3:
raise ValueError("Convention must have 3 letters.")
if convention[1] in (convention[0], convention[2]):
raise ValueError(f"Invalid convention {convention}.")
for letter in convention:
if letter not in ("X", "Y", "Z"):
raise ValueError(f"Invalid letter {letter} in convention string.")
matrices = map(_primary_matrix, convention, torch.unbind(euler_angles, -1))
return functools.reduce(torch.matmul, matrices)
def _angle_from_tan(
axis: str, other_axis: str, data, horizontal: bool, tait_bryan: bool
):
"""
Extract the first or third Euler angle from the two members of
the matrix which are positive constant times its sine and cosine.
Args:
axis: Axis label "X" or "Y or "Z" for the angle we are finding.
other_axis: Axis label "X" or "Y or "Z" for the middle axis in the
convention.
data: Rotation matrices as tensor of shape (..., 3, 3).
horizontal: Whether we are looking for the angle for the third axis,
which means the relevant entries are in the same row of the
rotation matrix. If not, they are in the same column.
tait_bryan: Whether the first and third axes in the convention differ.
Returns:
Euler Angles in radians for each matrix in data as a tensor
of shape (...).
"""
i1, i2 = {"X": (2, 1), "Y": (0, 2), "Z": (1, 0)}[axis]
if horizontal:
i2, i1 = i1, i2
even = (axis + other_axis) in ["XY", "YZ", "ZX"]
if horizontal == even:
return torch.atan2(data[..., i1], data[..., i2])
if tait_bryan:
return torch.atan2(-data[..., i2], data[..., i1])
return torch.atan2(data[..., i2], -data[..., i1])
def _index_from_letter(letter: str):
if letter == "X":
return 0
if letter == "Y":
return 1
if letter == "Z":
return 2
def matrix_to_euler_angles(matrix, convention: str):
"""
Convert rotations given as rotation matrices to Euler angles in radians.
Args:
matrix: Rotation matrices as tensor of shape (..., 3, 3).
convention: Convention string of three uppercase letters.
Returns:
Euler angles in radians as tensor of shape (..., 3).
"""
if len(convention) != 3:
raise ValueError("Convention must have 3 letters.")
if convention[1] in (convention[0], convention[2]):
raise ValueError(f"Invalid convention {convention}.")
for letter in convention:
if letter not in ("X", "Y", "Z"):
raise ValueError(f"Invalid letter {letter} in convention string.")
if matrix.size(-1) != 3 or matrix.size(-2) != 3:
raise ValueError(f"Invalid rotation matrix shape f{matrix.shape}.")
i0 = _index_from_letter(convention[0])
i2 = _index_from_letter(convention[2])
tait_bryan = i0 != i2
if tait_bryan:
central_angle = torch.asin(
matrix[..., i0, i2] * (-1.0 if i0 - i2 in [-1, 2] else 1.0)
)
else:
central_angle = torch.acos(matrix[..., i0, i0])
o = (
_angle_from_tan(
convention[0], convention[1], matrix[..., i2], False, tait_bryan
),
central_angle,
_angle_from_tan(
convention[2], convention[1], matrix[..., i0, :], True, tait_bryan
),
)
return torch.stack(o, -1)
def random_quaternions(
n: int, dtype: torch.dtype = None, device=None, requires_grad=False
):
"""
Generate random quaternions representing rotations,
i.e. versors with nonnegative real part.
Args:
n: Number to return.
dtype: Type to return.
device: Desired device of returned tensor. Default:
uses the current device for the default tensor type.
requires_grad: Whether the resulting tensor should have the gradient
flag set.
Returns:
Quaternions as tensor of shape (N, 4).
"""
o = torch.randn(
(n, 4), dtype=dtype, device=device, requires_grad=requires_grad
)
s = (o * o).sum(1)
o = o / _copysign(torch.sqrt(s), o[:, 0])[:, None]
return o
def random_rotations(
n: int, dtype: torch.dtype = None, device=None, requires_grad=False
):
"""
Generate random rotations as 3x3 rotation matrices.
Args:
n: Number to return.
dtype: Type to return.
device: Device of returned tensor. Default: if None,
uses the current device for the default tensor type.
requires_grad: Whether the resulting tensor should have the gradient
flag set.
Returns:
Rotation matrices as tensor of shape (n, 3, 3).
"""
quaternions = random_quaternions(
n, dtype=dtype, device=device, requires_grad=requires_grad
)
return quaternion_to_matrix(quaternions)
def random_rotation(
dtype: torch.dtype = None, device=None, requires_grad=False
):
"""
Generate a single random 3x3 rotation matrix.
Args:
dtype: Type to return
device: Device of returned tensor. Default: if None,
uses the current device for the default tensor type
requires_grad: Whether the resulting tensor should have the gradient
flag set
Returns:
Rotation matrix as tensor of shape (3, 3).
"""
return random_rotations(1, dtype, device, requires_grad)[0]
def standardize_quaternion(quaternions):
"""
Convert a unit quaternion to a standard form: one in which the real
part is non negative.
Args:
quaternions: Quaternions with real part first,
as tensor of shape (..., 4).
Returns:
Standardized quaternions as tensor of shape (..., 4).
"""
return torch.where(quaternions[..., 0:1] < 0, -quaternions, quaternions)
def quaternion_raw_multiply(a, b):
"""
Multiply two quaternions.
Usual torch rules for broadcasting apply.
Args:
a: Quaternions as tensor of shape (..., 4), real part first.
b: Quaternions as tensor of shape (..., 4), real part first.
Returns:
The product of a and b, a tensor of quaternions shape (..., 4).
"""
aw, ax, ay, az = torch.unbind(a, -1)
bw, bx, by, bz = torch.unbind(b, -1)
ow = aw * bw - ax * bx - ay * by - az * bz
ox = aw * bx + ax * bw + ay * bz - az * by
oy = aw * by - ax * bz + ay * bw + az * bx
oz = aw * bz + ax * by - ay * bx + az * bw
return torch.stack((ow, ox, oy, oz), -1)
def quaternion_multiply(a, b):
"""
Multiply two quaternions representing rotations, returning the quaternion
representing their composition, i.e. the versor with nonnegative real part.
Usual torch rules for broadcasting apply.
Args:
a: Quaternions as tensor of shape (..., 4), real part first.
b: Quaternions as tensor of shape (..., 4), real part first.
Returns:
The product of a and b, a tensor of quaternions of shape (..., 4).
"""
ab = quaternion_raw_multiply(a, b)
return standardize_quaternion(ab)
def quaternion_invert(quaternion):
"""
Given a quaternion representing rotation, get the quaternion representing
its inverse.
Args:
quaternion: Quaternions as tensor of shape (..., 4), with real part
first, which must be versors (unit quaternions).
Returns:
The inverse, a tensor of quaternions of shape (..., 4).
"""
return quaternion * quaternion.new_tensor([1, -1, -1, -1])
def quaternion_apply(quaternion, point):
"""
Apply the rotation given by a quaternion to a 3D point.
Usual torch rules for broadcasting apply.
Args:
quaternion: Tensor of quaternions, real part first, of shape (..., 4).
point: Tensor of 3D points of shape (..., 3).
Returns:
Tensor of rotated points of shape (..., 3).
"""
if point.size(-1) != 3:
raise ValueError(f"Points are not in 3D, f{point.shape}.")
real_parts = point.new_zeros(point.shape[:-1] + (1,))
point_as_quaternion = torch.cat((real_parts, point), -1)
out = quaternion_raw_multiply(
quaternion_raw_multiply(quaternion, point_as_quaternion),
quaternion_invert(quaternion),
)
return out[..., 1:]
pytorch3d/transforms/so3.py 0 → 100644
View file @ dbf06b50
#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
import torch
HAT_INV_SKEW_SYMMETRIC_TOL = 1e-5
def so3_relative_angle(R1, R2, cos_angle: bool = False):
"""
Calculates the relative angle (in radians) between pairs of
rotation matrices `R1` and `R2` with `angle = acos(0.5 * Trace(R1 R2^T)-1)`
.. note::
This corresponds to a geodesic distance on the 3D manifold of rotation
matrices.
Args:
R1: Batch of rotation matrices of shape `(minibatch, 3, 3)`.
R2: Batch of rotation matrices of shape `(minibatch, 3, 3)`.
cos_angle: If==True return cosine of the relative angle rather than
the angle itself. This can avoid the unstable
calculation of `acos`.
Returns:
Corresponding rotation angles of shape `(minibatch,)`.
If `cos_angle==True`, returns the cosine of the angles.
Raises:
ValueError if `R1` or `R2` is of incorrect shape.
ValueError if `R1` or `R2` has an unexpected trace.
"""
R12 = torch.bmm(R1, R2.permute(0, 2, 1))
return so3_rotation_angle(R12, cos_angle=cos_angle)
def so3_rotation_angle(R, eps: float = 1e-4, cos_angle: bool = False):
"""
Calculates angles (in radians) of a batch of rotation matrices `R` with
`angle = acos(0.5 * (Trace(R)-1))`. The trace of the
input matrices is checked to be in the valid range `[-1-eps,3+eps]`.
The `eps` argument is a small constant that allows for small errors
caused by limited machine precision.
Args:
R: Batch of rotation matrices of shape `(minibatch, 3, 3)`.
eps: Tolerance for the valid trace check.
cos_angle: If==True return cosine of the rotation angles rather than
the angle itself. This can avoid the unstable
calculation of `acos`.
Returns:
Corresponding rotation angles of shape `(minibatch,)`.
If `cos_angle==True`, returns the cosine of the angles.
Raises:
ValueError if `R` is of incorrect shape.
ValueError if `R` has an unexpected trace.
"""
N, dim1, dim2 = R.shape
if dim1 != 3 or dim2 != 3:
raise ValueError("Input has to be a batch of 3x3 Tensors.")
rot_trace = R[:, 0, 0] + R[:, 1, 1] + R[:, 2, 2]
if ((rot_trace < -1.0 - eps) + (rot_trace > 3.0 + eps)).any():
raise ValueError(
"A matrix has trace outside valid range [-1-eps,3+eps]."
)
# clamp to valid range
rot_trace = torch.clamp(rot_trace, -1.0, 3.0)
# phi ... rotation angle
phi = 0.5 * (rot_trace - 1.0)
if cos_angle:
return phi
else:
return phi.acos()
def so3_exponential_map(log_rot, eps: float = 0.0001):
"""
Convert a batch of logarithmic representations of rotation matrices `log_rot`
to a batch of 3x3 rotation matrices using Rodrigues formula [1].
In the logarithmic representation, each rotation matrix is represented as
a 3-dimensional vector (`log_rot`) who's l2-norm and direction correspond
to the magnitude of the rotation angle and the axis of rotation respectively.
The conversion has a singularity around `log(R) = 0`
which is handled by clamping controlled with the `eps` argument.
Args:
log_rot: Batch of vectors of shape `(minibatch , 3)`.
eps: A float constant handling the conversion singularity.
Returns:
Batch of rotation matrices of shape `(minibatch , 3 , 3)`.
Raises:
ValueError if `log_rot` is of incorrect shape.
[1] https://en.wikipedia.org/wiki/Rodrigues%27_rotation_formula
"""
_, dim = log_rot.shape
if dim != 3:
raise ValueError("Input tensor shape has to be Nx3.")
nrms = (log_rot * log_rot).sum(1)
# phis ... rotation angles
rot_angles = torch.clamp(nrms, eps).sqrt()
rot_angles_inv = 1.0 / rot_angles
fac1 = rot_angles_inv * rot_angles.sin()
fac2 = rot_angles_inv * rot_angles_inv * (1.0 - rot_angles.cos())
skews = hat(log_rot)
R = (
fac1[:, None, None] * skews
+ fac2[:, None, None] * torch.bmm(skews, skews)
+ torch.eye(3, dtype=log_rot.dtype, device=log_rot.device)[None]
)
return R
def so3_log_map(R, eps: float = 0.0001):
"""
Convert a batch of 3x3 rotation matrices `R`
to a batch of 3-dimensional matrix logarithms of rotation matrices
The conversion has a singularity around `(R=I)` which is handled
by clamping controlled with the `eps` argument.
Args:
R: batch of rotation matrices of shape `(minibatch, 3, 3)`.
eps: A float constant handling the conversion singularity.
Returns:
Batch of logarithms of input rotation matrices
of shape `(minibatch, 3)`.
Raises:
ValueError if `R` is of incorrect shape.
ValueError if `R` has an unexpected trace.
"""
N, dim1, dim2 = R.shape
if dim1 != 3 or dim2 != 3:
raise ValueError("Input has to be a batch of 3x3 Tensors.")
phi = so3_rotation_angle(R)
phi_valid = torch.clamp(phi.abs(), eps) * phi.sign()
log_rot_hat = (phi_valid / (2.0 * phi_valid.sin()))[:, None, None] * (
R - R.permute(0, 2, 1)
)
log_rot = hat_inv(log_rot_hat)
return log_rot
def hat_inv(h):
"""
Compute the inverse Hat operator [1] of a batch of 3x3 matrices.
Args:
h: Batch of skew-symmetric matrices of shape `(minibatch, 3, 3)`.
Returns:
Batch of 3d vectors of shape `(minibatch, 3, 3)`.
Raises:
ValueError if `h` is of incorrect shape.
ValueError if `h` not skew-symmetric.
[1] https://en.wikipedia.org/wiki/Hat_operator
"""
N, dim1, dim2 = h.shape
if dim1 != 3 or dim2 != 3:
raise ValueError("Input has to be a batch of 3x3 Tensors.")
ss_diff = (h + h.permute(0, 2, 1)).abs().max()
if float(ss_diff) > HAT_INV_SKEW_SYMMETRIC_TOL:
raise ValueError("One of input matrices not skew-symmetric.")
x = h[:, 2, 1]
y = h[:, 0, 2]
z = h[:, 1, 0]
v = torch.stack((x, y, z), dim=1)
return v
def hat(v):
"""
Compute the Hat operator [1] of a batch of 3D vectors.
Args:
v: Batch of vectors of shape `(minibatch , 3)`.
Returns:
Batch of skew-symmetric matrices of shape
`(minibatch, 3 , 3)` where each matrix is of the form:
`[ 0 -v_z v_y ]
[ v_z 0 -v_x ]
[ -v_y v_x 0 ]`
Raises:
ValueError if `v` is of incorrect shape.
[1] https://en.wikipedia.org/wiki/Hat_operator
"""
N, dim = v.shape
if dim != 3:
raise ValueError("Input vectors have to be 3-dimensional.")
h = v.new_zeros(N, 3, 3)
x, y, z = v.unbind(1)
h[:, 0, 1] = -z
h[:, 0, 2] = y
h[:, 1, 0] = z
h[:, 1, 2] = -x
h[:, 2, 0] = -y
h[:, 2, 1] = x
return h
pytorch3d/transforms/transform3d.py 0 → 100644
View file @ dbf06b50
#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
import math
import torch
class Transform3d:
"""
A Transform3d object encapsulates a batch of N 3D transformations, and knows
how to transform points and normal vectors. Suppose that t is a Transform3d;
then we can do the following:
.. code-block:: python
N = len(t)
points = torch.randn(N, P, 3)
normals = torch.randn(N, P, 3)
points_transformed = t.transform_points(points) # => (N, P, 3)
normals_transformed = t.transform_points(normals) # => (N, P, 3)
BROADCASTING
Transform3d objects supports broadcasting. Suppose that t1 and tN are
Transform3D objects with len(t1) == 1 and len(tN) == N respectively. Then we
can broadcast transforms like this:
.. code-block:: python
t1.transform_points(torch.randn(P, 3)) # => (P, 3)
t1.transform_points(torch.randn(1, P, 3)) # => (1, P, 3)
t1.transform_points(torch.randn(M, P, 3)) # => (M, P, 3)
tN.transform_points(torch.randn(P, 3)) # => (N, P, 3)
tN.transform_points(torch.randn(1, P, 3)) # => (N, P, 3)
COMBINING TRANSFORMS
Transform3d objects can be combined in two ways: composing and stacking.
Composing is function composition. Given Transform3d objects t1, t2, t3,
the following all compute the same thing:
.. code-block:: python
y1 = t3.transform_points(t2.transform_points(t2.transform_points(x)))
y2 = t1.compose(t2).compose(t3).transform_points()
y3 = t1.compose(t2, t3).transform_points()
Composing transforms should broadcast.
.. code-block:: python
if len(t1) == 1 and len(t2) == N, then len(t1.compose(t2)) == N.
We can also stack a sequence of Transform3d objects, which represents
composition along the batch dimension; then the following should compute the
same thing.
.. code-block:: python
N, M = len(tN), len(tM)
xN = torch.randn(N, P, 3)
xM = torch.randn(M, P, 3)
y1 = torch.cat([tN.transform_points(xN), tM.transform_points(xM)], dim=0)
y2 = tN.stack(tM).transform_points(torch.cat([xN, xM], dim=0))
BUILDING TRANSFORMS
We provide convenience methods for easily building Transform3d objects
as compositions of basic transforms.
.. code-block:: python
# Scale by 0.5, then translate by (1, 2, 3)
t1 = Transform3d().scale(0.5).translate(1, 2, 3)
# Scale each axis by a different amount, then translate, then scale
t2 = Transform3d().scale(1, 3, 3).translate(2, 3, 1).scale(2.0)
t3 = t1.compose(t2)
tN = t1.stack(t3, t3)
BACKPROP THROUGH TRANSFORMS
When building transforms, we can also parameterize them by Torch tensors;
in this case we can backprop through the construction and application of
Transform objects, so they could be learned via gradient descent or
predicted by a neural network.
.. code-block:: python
s1_params = torch.randn(N, requires_grad=True)
t_params = torch.randn(N, 3, requires_grad=True)
s2_params = torch.randn(N, 3, requires_grad=True)
t = Transform3d().scale(s1_params).translate(t_params).scale(s2_params)
x = torch.randn(N, 3)
y = t.transform_points(x)
loss = compute_loss(y)
loss.backward()
with torch.no_grad():
s1_params -= lr * s1_params.grad
t_params -= lr * t_params.grad
s2_params -= lr * s2_params.grad
"""
def __init__(self, dtype=torch.float32, device="cpu"):
"""
This class assumes a row major ordering for all matrices.
"""
self._matrix = torch.eye(4, dtype=dtype, device=device).view(1, 4, 4)
self._transforms = [] # store transforms to compose
self._lu = None
self.device = device
def __len__(self):
return self.get_matrix().shape[0]
def compose(self, *others):
"""
Return a new Transform3d with the tranforms to compose stored as
an internal list.
Args:
*others: Any number of Transform3d objects
Returns:
A new Transform3d with the stored transforms
"""
out = Transform3d(device=self.device)
out._matrix = self._matrix.clone()
for other in others:
if not isinstance(other, Transform3d):
msg = "Only possible to compose Transform3d objects; got %s"
raise ValueError(msg % type(other))
out._transforms = self._transforms + list(others)
return out
def get_matrix(self):
"""
Return a matrix which is the result of composing this transform
with others stored in self.transforms. Where necessary transforms
are broadcast against each other.
For example, if self.transforms contains transforms t1, t2, and t3, and
given a set of points x, the following should be true:
.. code-block:: python
y1 = t1.compose(t2, t3).transform(x)
y2 = t3.transform(t2.transform(t1.transform(x)))
y1.get_matrix() == y2.get_matrix()
Returns:
A transformation matrix representing the composed inputs.
"""
composed_matrix = self._matrix.clone()
if len(self._transforms) > 0:
for other in self._transforms:
other_matrix = other.get_matrix()
composed_matrix = _broadcast_bmm(composed_matrix, other_matrix)
return composed_matrix
def _get_matrix_inverse(self):
"""
Return the inverse of self._matrix.
"""
return torch.inverse(self._matrix)
def inverse(self, invert_composed: bool = False):
"""
Returns a new Transform3D object that represents an inverse of the
current transformation.
Args:
invert_composed:
- True: First compose the list of stored transformations
and then apply inverse to the result. This is
potentially slower for classes of transformations
with inverses that can be computed efficiently
(e.g. rotations and translations).
- False: Invert the individual stored transformations
independently without composing them.
Returns:
A new Transform3D object contaning the inverse of the original
transformation.
"""
tinv = Transform3d(device=self.device)
if invert_composed:
# first compose then invert
tinv._matrix = torch.inverse(self.get_matrix())
else:
# self._get_matrix_inverse() implements efficient inverse
# of self._matrix
i_matrix = self._get_matrix_inverse()
# 2 cases:
if len(self._transforms) > 0:
# a) Either we have a non-empty list of transforms:
# Here we take self._matrix and append its inverse at the
# end of the reverted _transforms list. After composing
# the transformations with get_matrix(), this correctly
# right-multiplies by the inverse of self._matrix
# at the end of the composition.
tinv._transforms = [
t.inverse() for t in reversed(self._transforms)
]
last = Transform3d(device=self.device)
last._matrix = i_matrix
tinv._transforms.append(last)
else:
# b) Or there are no stored transformations
# we just set inverted matrix
tinv._matrix = i_matrix
return tinv
def stack(self, *others):
transforms = [self] + list(others)
matrix = torch.cat([t._matrix for t in transforms], dim=0)
out = Transform3d()
out._matrix = matrix
return out
def transform_points(self, points, eps: float = None):
"""
Use this transform to transform a set of 3D points. Assumes row major
ordering of the input points.
Args:
points: Tensor of shape (P, 3) or (N, P, 3)
eps: If eps!=None, the argument is used to clamp the
last coordinate before peforming the final division.
The clamping corresponds to:
last_coord := (last_coord.sign() + (last_coord==0)) *
torch.clamp(last_coord.abs(), eps),
i.e. the last coordinates that are exactly 0 will
be clamped to +eps.
Returns:
points_out: points of shape (N, P, 3) or (P, 3) depending
on the dimensions of the transform
"""
points_batch = points.clone()
if points_batch.dim() == 2:
points_batch = points_batch[None] # (P, 3) -> (1, P, 3)
if points_batch.dim() != 3:
msg = "Expected points to have dim = 2 or dim = 3: got shape %r"
raise ValueError(msg % points.shape)
N, P, _3 = points_batch.shape
ones = torch.ones(N, P, 1, dtype=points.dtype, device=points.device)
points_batch = torch.cat([points_batch, ones], dim=2)
composed_matrix = self.get_matrix()
points_out = _broadcast_bmm(points_batch, composed_matrix)
denom = points_out[..., 3:] # denominator
if eps is not None:
denom_sign = denom.sign() + (denom == 0.0).type_as(denom)
denom = denom_sign * torch.clamp(denom.abs(), eps)
points_out = points_out[..., :3] / denom
# When transform is (1, 4, 4) and points is (P, 3) return
# points_out of shape (P, 3)
if points_out.shape[0] == 1 and points.dim() == 2:
points_out = points_out.reshape(points.shape)
return points_out
def transform_normals(self, normals):
"""
Use this transform to transform a set of normal vectors.
Args:
normals: Tensor of shape (P, 3) or (N, P, 3)
Returns:
normals_out: Tensor of shape (P, 3) or (N, P, 3) depending
on the dimensions of the transform
"""
if normals.dim() not in [2, 3]:
msg = "Expected normals to have dim = 2 or dim = 3: got shape %r"
raise ValueError(msg % normals.shape)
composed_matrix = self.get_matrix()
# TODO: inverse is bad! Solve a linear system instead
mat = composed_matrix[:, :3, :3]
normals_out = _broadcast_bmm(normals, mat.transpose(1, 2).inverse())
# This doesn't pass unit tests. TODO investigate further
# if self._lu is None:
# self._lu = self._matrix[:, :3, :3].transpose(1, 2).lu()
# normals_out = normals.lu_solve(*self._lu)
# When transform is (1, 4, 4) and normals is (P, 3) return
# normals_out of shape (P, 3)
if normals_out.shape[0] == 1 and normals.dim() == 2:
normals_out = normals_out.reshape(normals.shape)
return normals_out
def translate(self, *args, **kwargs):
return self.compose(Translate(device=self.device, *args, **kwargs))
def scale(self, *args, **kwargs):
return self.compose(Scale(device=self.device, *args, **kwargs))
def rotate_axis_angle(self, *args, **kwargs):
return self.compose(
RotateAxisAngle(device=self.device, *args, **kwargs)
)
def clone(self):
"""
Deep copy of Transforms object. All internal tensors are cloned
individually.
Returns:
new Transforms object.
"""
other = Transform3d(device=self.device)
if self._lu is not None:
other._lu = [l.clone() for l in self._lu]
other._matrix = self._matrix.clone()
other._transforms = [t.clone() for t in self._transforms]
return other
def to(self, device, copy: bool = False, dtype=None):
"""
Match functionality of torch.Tensor.to()
If copy = True or the self Tensor is on a different device, the
returned tensor is a copy of self with the desired torch.device.
If copy = False and the self Tensor already has the correct torch.device,
then self is returned.
Args:
device: Device id for the new tensor.
copy: Boolean indicator whether or not to clone self. Default False.
dtype: If not None, casts the internal tensor variables
to a given torch.dtype.
Returns:
Transform3d object.
"""
if not copy and self.device == device:
return self
other = self.clone()
if self.device != device:
other.device = device
other._matrix = self._matrix.to(device=device, dtype=dtype)
for t in other._transforms:
t.to(device, copy=copy, dtype=dtype)
return other
def cpu(self):
return self.to(torch.device("cpu"))
def cuda(self):
return self.to(torch.device("cuda"))
class Translate(Transform3d):
def __init__(
self, x, y=None, z=None, dtype=torch.float32, device: str = "cpu"
):
"""
Create a new Transform3d representing 3D translations.
Option I: Translate(xyz, dtype=torch.float32, device='cpu')
xyz should be a tensor of shape (N, 3)
Option II: Translate(x, y, z, dtype=torch.float32, device='cpu')
Here x, y, and z will be broadcast against each other and
concatenated to form the translation. Each can be:
- A python scalar
- A torch scalar
- A 1D torch tensor
"""
super().__init__(device=device)
xyz = _handle_input(x, y, z, dtype, device, "Translate")
N = xyz.shape[0]
mat = torch.eye(4, dtype=dtype, device=device)
mat = mat.view(1, 4, 4).repeat(N, 1, 1)
mat[:, 3, :3] = xyz
self._matrix = mat
def _get_matrix_inverse(self):
"""
Return the inverse of self._matrix.
"""
inv_mask = self._matrix.new_ones([1, 4, 4])
inv_mask[0, 3, :3] = -1.0
i_matrix = self._matrix * inv_mask
return i_matrix
class Scale(Transform3d):
def __init__(
self, x, y=None, z=None, dtype=torch.float32, device: str = "cpu"
):
"""
A Transform3d representing a scaling operation, with different scale
factors along each coordinate axis.
Option I: Scale(s, dtype=torch.float32, device='cpu')
s can be one of
- Python scalar or torch scalar: Single uniform scale
- 1D torch tensor of shape (N,): A batch of uniform scale
- 2D torch tensor of shape (N, 3): Scale differently along each axis
Option II: Scale(x, y, z, dtype=torch.float32, device='cpu')
Each of x, y, and z can be one of
- python scalar
- torch scalar
- 1D torch tensor
"""
super().__init__(device=device)
xyz = _handle_input(
x, y, z, dtype, device, "scale", allow_singleton=True
)
N = xyz.shape[0]
# TODO: Can we do this all in one go somehow?
mat = torch.eye(4, dtype=dtype, device=device)
mat = mat.view(1, 4, 4).repeat(N, 1, 1)
mat[:, 0, 0] = xyz[:, 0]
mat[:, 1, 1] = xyz[:, 1]
mat[:, 2, 2] = xyz[:, 2]
self._matrix = mat
def _get_matrix_inverse(self):
"""
Return the inverse of self._matrix.
"""
xyz = torch.stack([self._matrix[:, i, i] for i in range(4)], dim=1)
ixyz = 1.0 / xyz
imat = torch.diag_embed(ixyz, dim1=1, dim2=2)
return imat
class Rotate(Transform3d):
def __init__(
self,
R,
dtype=torch.float32,
device: str = "cpu",
orthogonal_tol: float = 1e-5,
):
"""
Create a new Transform3d representing 3D rotation using a rotation
matrix as the input.
Args:
R: a tensor of shape (3, 3) or (N, 3, 3)
orthogonal_tol: tolerance for the test of the orthogonality of R
"""
super().__init__(device=device)
if R.dim() == 2:
R = R[None]
if R.shape[-2:] != (3, 3):
msg = "R must have shape (3, 3) or (N, 3, 3); got %s"
raise ValueError(msg % repr(R.shape))
R = R.to(dtype=dtype).to(device=device)
_check_valid_rotation_matrix(R, tol=orthogonal_tol)
N = R.shape[0]
mat = torch.eye(4, dtype=dtype, device=device)
mat = mat.view(1, 4, 4).repeat(N, 1, 1)
mat[:, :3, :3] = R
self._matrix = mat
def _get_matrix_inverse(self):
"""
Return the inverse of self._matrix.
"""
return self._matrix.permute(0, 2, 1).contiguous()
class RotateAxisAngle(Rotate):
def __init__(
self,
angle,
axis: str = "X",
degrees: bool = True,
dtype=torch.float64,
device: str = "cpu",
):
"""
Create a new Transform3d representing 3D rotation about an axis
by an angle.
Args:
angle:
- A torch tensor of shape (N, 1)
- A python scalar
- A torch scalar
axis:
string: one of ["X", "Y", "Z"] indicating the axis about which
to rotate.
NOTE: All batch elements are rotated about the same axis.
"""
axis = axis.upper()
if axis not in ["X", "Y", "Z"]:
msg = "Expected axis to be one of ['X', 'Y', 'Z']; got %s"
raise ValueError(msg % axis)
angle = _handle_angle_input(angle, dtype, device, "RotateAxisAngle")
angle = (angle / 180.0 * math.pi) if degrees else angle
N = angle.shape[0]
cos = torch.cos(angle)
sin = torch.sin(angle)
one = torch.ones_like(angle)
zero = torch.zeros_like(angle)
if axis == "X":
R_flat = (one, zero, zero, zero, cos, -sin, zero, sin, cos)
if axis == "Y":
R_flat = (cos, zero, sin, zero, one, zero, -sin, zero, cos)
if axis == "Z":
R_flat = (cos, -sin, zero, sin, cos, zero, zero, zero, one)
R = torch.stack(R_flat, -1).reshape((N, 3, 3))
super().__init__(device=device, R=R)
def _handle_coord(c, dtype, device):
"""
Helper function for _handle_input.
Args:
c: Python scalar, torch scalar, or 1D torch tensor
Returns:
c_vec: 1D torch tensor
"""
if not torch.is_tensor(c):
c = torch.tensor(c, dtype=dtype, device=device)
if c.dim() == 0:
c = c.view(1)
return c
def _handle_input(
x, y, z, dtype, device, name: str, allow_singleton: bool = False
):
"""
Helper function to handle parsing logic for building transforms. The output
is always a tensor of shape (N, 3), but there are several types of allowed
input.
Case I: Single Matrix
In this case x is a tensor of shape (N, 3), and y and z are None. Here just
return x.
Case II: Vectors and Scalars
In this case each of x, y, and z can be one of the following
- Python scalar
- Torch scalar
- Torch tensor of shape (N, 1) or (1, 1)
In this case x, y and z are broadcast to tensors of shape (N, 1)
and concatenated to a tensor of shape (N, 3)
Case III: Singleton (only if allow_singleton=True)
In this case y and z are None, and x can be one of the following:
- Python scalar
- Torch scalar
- Torch tensor of shape (N, 1) or (1, 1)
Here x will be duplicated 3 times, and we return a tensor of shape (N, 3)
Returns:
xyz: Tensor of shape (N, 3)
"""
# If x is actually a tensor of shape (N, 3) then just return it
if torch.is_tensor(x) and x.dim() == 2:
if x.shape[1] != 3:
msg = "Expected tensor of shape (N, 3); got %r (in %s)"
raise ValueError(msg % (x.shape, name))
if y is not None or z is not None:
msg = "Expected y and z to be None (in %s)" % name
raise ValueError(msg)
return x
if allow_singleton and y is None and z is None:
y = x
z = x
# Convert all to 1D tensors
xyz = [_handle_coord(c, dtype, device) for c in [x, y, z]]
# Broadcast and concatenate
sizes = [c.shape[0] for c in xyz]
N = max(sizes)
for c in xyz:
if c.shape[0] != 1 and c.shape[0] != N:
msg = "Got non-broadcastable sizes %r (in %s)" % (sizes, name)
raise ValueError(msg)
xyz = [c.expand(N) for c in xyz]
xyz = torch.stack(xyz, dim=1)
return xyz
def _handle_angle_input(x, dtype, device: str, name: str):
"""
Helper function for building a rotation function using angles.
The output is always of shape (N, 1).
The input can be one of:
- Torch tensor (N, 1) or (N)
- Python scalar
- Torch scalar
"""
# If x is actually a tensor of shape (N, 1) then just return it
if torch.is_tensor(x) and x.dim() == 2:
if x.shape[1] != 1:
msg = "Expected tensor of shape (N, 1); got %r (in %s)"
raise ValueError(msg % (x.shape, name))
return x
else:
return _handle_coord(x, dtype, device)
def _broadcast_bmm(a, b):
"""
Batch multiply two matrices and broadcast if necessary.
Args:
a: torch tensor of shape (P, K) or (M, P, K)
b: torch tensor of shape (N, K, K)
Returns:
a and b broadcast multipled. The output batch dimension is max(N, M).
To broadcast transforms across a batch dimension if M != N then
expect that either M = 1 or N = 1. The tensor with batch dimension 1 is
expanded to have shape N or M.
"""
if a.dim() == 2:
a = a[None]
if len(a) != len(b):
if not ((len(a) == 1) or (len(b) == 1)):
msg = "Expected batch dim for bmm to be equal or 1; got %r, %r"
raise ValueError(msg % (a.shape, b.shape))
if len(a) == 1:
a = a.expand(len(b), -1, -1)
if len(b) == 1:
b = b.expand(len(a), -1, -1)
return a.bmm(b)
def _check_valid_rotation_matrix(R, tol: float = 1e-7):
"""
Determine if R is a valid rotation matrix by checking it satisfies the
following conditions:
``RR^T = I and det(R) = 1``
Args:
R: an (N, 3, 3) matrix
Returns:
None
Prints an warning if R is an invalid rotation matrix. Else return.
"""
N = R.shape[0]
eye = torch.eye(3, dtype=R.dtype, device=R.device)
eye = eye.view(1, 3, 3).expand(N, -1, -1)
orthogonal = torch.allclose(R.bmm(R.transpose(1, 2)), eye, atol=tol)
det_R = torch.det(R)
no_distortion = torch.allclose(det_R, torch.ones_like(det_R))
if not (orthogonal and no_distortion):
msg = "R is not a valid rotation matrix"
print(msg)
return
pytorch3d/utils/__init__.py 0 → 100644
View file @ dbf06b50
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
from .ico_sphere import ico_sphere
__all__ = [k for k in globals().keys() if not k.startswith("_")]
pytorch3d/utils/ico_sphere.py 0 → 100644
View file @ dbf06b50
#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
import torch
from pytorch3d.ops.subdivide_meshes import SubdivideMeshes
from pytorch3d.structures.meshes import Meshes
# Vertex coordinates for a level 0 ico-sphere.
_ico_verts0 = [
[-0.5257, 0.8507, 0.0000],
[0.5257, 0.8507, 0.0000],
[-0.5257, -0.8507, 0.0000],
[0.5257, -0.8507, 0.0000],
[0.0000, -0.5257, 0.8507],
[0.0000, 0.5257, 0.8507],
[0.0000, -0.5257, -0.8507],
[0.0000, 0.5257, -0.8507],
[0.8507, 0.0000, -0.5257],
[0.8507, 0.0000, 0.5257],
[-0.8507, 0.0000, -0.5257],
[-0.8507, 0.0000, 0.5257],
]
# Faces for level 0 ico-sphere
_ico_faces0 = [
[0, 11, 5],
[0, 5, 1],
[0, 1, 7],
[0, 7, 10],
[0, 10, 11],
[1, 5, 9],
[5, 11, 4],
[11, 10, 2],
[10, 7, 6],
[7, 1, 8],
[3, 9, 4],
[3, 4, 2],
[3, 2, 6],
[3, 6, 8],
[3, 8, 9],
[4, 9, 5],
[2, 4, 11],
[6, 2, 10],
[8, 6, 7],
[9, 8, 1],
]
def ico_sphere(level: int = 0, device=None):
"""
Create verts and faces for a unit ico-sphere, with all faces oriented
consistently.
Args:
level: integer specifying the number of iterations for subdivision
of the mesh faces. Each additional level will result in four new
faces per face.
device: A torch.device object on which the outputs will be allocated.
Returns:
Meshes object with verts and faces.
"""
if device is None:
device = torch.device("cpu")
if level < 0:
raise ValueError("level must be >= 0.")
if level == 0:
verts = torch.tensor(_ico_verts0, dtype=torch.float32, device=device)
faces = torch.tensor(_ico_faces0, dtype=torch.int64, device=device)
else:
mesh = ico_sphere(level - 1, device)
subdivide = SubdivideMeshes()
mesh = subdivide(mesh)
verts = mesh.verts_list()[0]
verts /= verts.norm(p=2, dim=1, keepdim=True)
faces = mesh.faces_list()[0]
return Meshes(verts=[verts], faces=[faces])
setup.cfg 0 → 100644
View file @ dbf06b50
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
[isort]
line_length=80
include_trailing_comma=True
multi_line_output=3
known_standard_library=numpy,setuptools
known_myself=pytorch3d
known_third_party=fvcore,torch,torchvision,matplotlib,mpl_toolkits,PIL,yaml,jinja2,requests
no_lines_before=STDLIB,THIRDPARTY
sections=FUTURE,STDLIB,THIRDPARTY,myself,FIRSTPARTY,LOCALFOLDER
default_section=FIRSTPARTY
setup.py 0 → 100644
View file @ dbf06b50
#!/usr/bin/env python
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
import glob
import os
from setuptools import find_packages, setup
import torch
from torch.utils.cpp_extension import CUDA_HOME, CppExtension, CUDAExtension
def get_extensions():
this_dir = os.path.dirname(os.path.abspath(__file__))
extensions_dir = os.path.join(this_dir, "pytorch3d", "csrc")
main_source = os.path.join(extensions_dir, "ext.cpp")
sources = glob.glob(os.path.join(extensions_dir, "**", "*.cpp"))
source_cuda = glob.glob(os.path.join(extensions_dir, "**", "*.cu"))
sources = [main_source] + sources
extension = CppExtension
extra_compile_args = {"cxx": ["-std=c++17"]}
define_macros = []
if torch.cuda.is_available() and CUDA_HOME is not None:
extension = CUDAExtension
sources += source_cuda
define_macros += [("WITH_CUDA", None)]
extra_compile_args["nvcc"] = [
"-DCUDA_HAS_FP16=1",
"-D__CUDA_NO_HALF_OPERATORS__",
"-D__CUDA_NO_HALF_CONVERSIONS__",
"-D__CUDA_NO_HALF2_OPERATORS__",
]
# It's better if pytorch can do this by default ..
CC = os.environ.get("CC", None)
if CC is not None:
extra_compile_args["nvcc"].append("-ccbin={}".format(CC))
sources = [os.path.join(extensions_dir, s) for s in sources]
include_dirs = [extensions_dir]
ext_modules = [
extension(
"pytorch3d._C",
sources,
include_dirs=include_dirs,
define_macros=define_macros,
extra_compile_args=extra_compile_args,
)
]
return ext_modules
setup(
name="pytorch3d",
version="0.1",
author="FAIR",
url="https://github.com/facebookresearch/pytorch3d",
description="PyTorch3d is FAIR's library of reusable components "
"for deep Learning with 3D data.",
packages=find_packages(exclude=("configs", "tests")),
install_requires=["torchvision>=0.4", "fvcore"],
extras_require={
"all": ["matplotlib", "tqdm>4.29.0", "imageio", "ipywidgets"],
"dev": ["flake8", "isort", "black==19.3b0"],
},
ext_modules=get_extensions(),
cmdclass={"build_ext": torch.utils.cpp_extension.BuildExtension},
)
tests/__init__.py 0 → 100644
View file @ dbf06b50
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
tests/bm_chamfer.py 0 → 100644
View file @ dbf06b50
#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
import torch
from fvcore.common.benchmark import benchmark
from test_chamfer import TestChamfer
def bm_chamfer() -> None:
kwargs_list_naive = [
{"batch_size": 1, "P1": 32, "P2": 64, "return_normals": False},
{"batch_size": 1, "P1": 32, "P2": 64, "return_normals": True},
{"batch_size": 32, "P1": 32, "P2": 64, "return_normals": False},
]
benchmark(
TestChamfer.chamfer_naive_with_init,
"CHAMFER_NAIVE",
kwargs_list_naive,
warmup_iters=1,
)
if torch.cuda.is_available():
kwargs_list = kwargs_list_naive + [
{"batch_size": 1, "P1": 1000, "P2": 3000, "return_normals": False},
{"batch_size": 1, "P1": 1000, "P2": 30000, "return_normals": True},
]
benchmark(
TestChamfer.chamfer_with_init,
"CHAMFER",
kwargs_list,
warmup_iters=1,
)
tests/bm_cubify.py 0 → 100644
View file @ dbf06b50
#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
from fvcore.common.benchmark import benchmark
from test_cubify import TestCubify
def bm_cubify() -> None:
kwargs_list = [
{"batch_size": 32, "V": 16},
{"batch_size": 64, "V": 16},
{"batch_size": 16, "V": 32},
]
benchmark(
TestCubify.cubify_with_init, "CUBIFY", kwargs_list, warmup_iters=1
)
tests/bm_graph_conv.py 0 → 100644
View file @ dbf06b50
#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
from itertools import product
import torch
from fvcore.common.benchmark import benchmark
from test_graph_conv import TestGraphConv
def bm_graph_conv() -> None:
backends = ["cpu"]
if torch.cuda.is_available():
backends.append("cuda")
kwargs_list = []
gconv_dim = [128, 256]
num_meshes = [32, 64]
num_verts = [100]
num_faces = [1000]
directed = [False, True]
test_cases = product(
gconv_dim, num_meshes, num_verts, num_faces, directed, backends
)
for case in test_cases:
g, n, v, f, d, b = case
kwargs_list.append(
{
"gconv_dim": g,
"num_meshes": n,
"num_verts": v,
"num_faces": f,
"directed": d,
"backend": b,
}
)
benchmark(
TestGraphConv.graph_conv_forward_backward,
"GRAPH CONV",
kwargs_list,
warmup_iters=1,
)
tests/bm_main.py 0 → 100644
View file @ dbf06b50
#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
import glob
import importlib
from os.path import basename, dirname, isfile, join, sys
if __name__ == "__main__":
# pyre-ignore[16]
if len(sys.argv) > 1:
# Parse from flags.
# pyre-ignore[16]
module_names = [n for n in sys.argv if n.startswith("bm_")]
else:
# Get all the benchmark files (starting with "bm_").
bm_files = glob.glob(join(dirname(__file__), "bm_*.py"))
module_names = [
basename(f)[:-3]
for f in bm_files
if isfile(f) and not f.endswith("bm_main.py")
]
for module_name in module_names:
module = importlib.import_module(module_name)
for attr in dir(module):
# Run all the functions with names "bm_*" in the module.
if attr.startswith("bm_"):
print(
"Running benchmarks for " + module_name + "/" + attr + "..."
)
getattr(module, attr)()
tests/bm_mesh_edge_loss.py 0 → 100644
View file @ dbf06b50
#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
from itertools import product
from fvcore.common.benchmark import benchmark
from test_mesh_edge_loss import TestMeshEdgeLoss
def bm_mesh_edge_loss() -> None:
kwargs_list = []
num_meshes = [1, 16, 32]
max_v = [100, 10000]
max_f = [300, 30000]
test_cases = product(num_meshes, max_v, max_f)
for case in test_cases:
n, v, f = case
kwargs_list.append({"num_meshes": n, "max_v": v, "max_f": f})
benchmark(
TestMeshEdgeLoss.mesh_edge_loss,
"MESH_EDGE_LOSS",
kwargs_list,
warmup_iters=1,
)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment