|
|
|
|
|
|
|
|
|
import pickle
|
|
from functools import lru_cache
|
|
from typing import Dict, Optional, Tuple
|
|
import torch
|
|
|
|
from detectron2.utils.file_io import PathManager
|
|
|
|
from densepose.data.meshes.catalog import MeshCatalog, MeshInfo
|
|
|
|
|
|
def _maybe_copy_to_device(
|
|
attribute: Optional[torch.Tensor], device: torch.device
|
|
) -> Optional[torch.Tensor]:
|
|
if attribute is None:
|
|
return None
|
|
return attribute.to(device)
|
|
|
|
|
|
class Mesh:
|
|
def __init__(
|
|
self,
|
|
vertices: Optional[torch.Tensor] = None,
|
|
faces: Optional[torch.Tensor] = None,
|
|
geodists: Optional[torch.Tensor] = None,
|
|
symmetry: Optional[Dict[str, torch.Tensor]] = None,
|
|
texcoords: Optional[torch.Tensor] = None,
|
|
mesh_info: Optional[MeshInfo] = None,
|
|
device: Optional[torch.device] = None,
|
|
):
|
|
"""
|
|
Args:
|
|
vertices (tensor [N, 3] of float32): vertex coordinates in 3D
|
|
faces (tensor [M, 3] of long): triangular face represented as 3
|
|
vertex indices
|
|
geodists (tensor [N, N] of float32): geodesic distances from
|
|
vertex `i` to vertex `j` (optional, default: None)
|
|
symmetry (dict: str -> tensor): various mesh symmetry data:
|
|
- "vertex_transforms": vertex mapping under horizontal flip,
|
|
tensor of size [N] of type long; vertex `i` is mapped to
|
|
vertex `tensor[i]` (optional, default: None)
|
|
texcoords (tensor [N, 2] of float32): texture coordinates, i.e. global
|
|
and normalized mesh UVs (optional, default: None)
|
|
mesh_info (MeshInfo type): necessary to load the attributes on-the-go,
|
|
can be used instead of passing all the variables one by one
|
|
device (torch.device): device of the Mesh. If not provided, will use
|
|
the device of the vertices
|
|
"""
|
|
self._vertices = vertices
|
|
self._faces = faces
|
|
self._geodists = geodists
|
|
self._symmetry = symmetry
|
|
self._texcoords = texcoords
|
|
self.mesh_info = mesh_info
|
|
self.device = device
|
|
|
|
assert self._vertices is not None or self.mesh_info is not None
|
|
|
|
all_fields = [self._vertices, self._faces, self._geodists, self._texcoords]
|
|
|
|
if self.device is None:
|
|
for field in all_fields:
|
|
if field is not None:
|
|
self.device = field.device
|
|
break
|
|
if self.device is None and symmetry is not None:
|
|
for key in symmetry:
|
|
self.device = symmetry[key].device
|
|
break
|
|
self.device = torch.device("cpu") if self.device is None else self.device
|
|
|
|
assert all([var.device == self.device for var in all_fields if var is not None])
|
|
if symmetry:
|
|
assert all(symmetry[key].device == self.device for key in symmetry)
|
|
if texcoords and vertices:
|
|
assert len(vertices) == len(texcoords)
|
|
|
|
def to(self, device: torch.device):
|
|
device_symmetry = self._symmetry
|
|
if device_symmetry:
|
|
device_symmetry = {key: value.to(device) for key, value in device_symmetry.items()}
|
|
return Mesh(
|
|
_maybe_copy_to_device(self._vertices, device),
|
|
_maybe_copy_to_device(self._faces, device),
|
|
_maybe_copy_to_device(self._geodists, device),
|
|
device_symmetry,
|
|
_maybe_copy_to_device(self._texcoords, device),
|
|
self.mesh_info,
|
|
device,
|
|
)
|
|
|
|
@property
|
|
def vertices(self):
|
|
if self._vertices is None and self.mesh_info is not None:
|
|
self._vertices = load_mesh_data(self.mesh_info.data, "vertices", self.device)
|
|
return self._vertices
|
|
|
|
@property
|
|
def faces(self):
|
|
if self._faces is None and self.mesh_info is not None:
|
|
self._faces = load_mesh_data(self.mesh_info.data, "faces", self.device)
|
|
return self._faces
|
|
|
|
@property
|
|
def geodists(self):
|
|
if self._geodists is None and self.mesh_info is not None:
|
|
self._geodists = load_mesh_auxiliary_data(self.mesh_info.geodists, self.device)
|
|
return self._geodists
|
|
|
|
@property
|
|
def symmetry(self):
|
|
if self._symmetry is None and self.mesh_info is not None:
|
|
self._symmetry = load_mesh_symmetry(self.mesh_info.symmetry, self.device)
|
|
return self._symmetry
|
|
|
|
@property
|
|
def texcoords(self):
|
|
if self._texcoords is None and self.mesh_info is not None:
|
|
self._texcoords = load_mesh_auxiliary_data(self.mesh_info.texcoords, self.device)
|
|
return self._texcoords
|
|
|
|
def get_geodists(self):
|
|
if self.geodists is None:
|
|
self.geodists = self._compute_geodists()
|
|
return self.geodists
|
|
|
|
def _compute_geodists(self):
|
|
|
|
geodists = None
|
|
return geodists
|
|
|
|
|
|
def load_mesh_data(
|
|
mesh_fpath: str, field: str, device: Optional[torch.device] = None
|
|
) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:
|
|
with PathManager.open(mesh_fpath, "rb") as hFile:
|
|
|
|
|
|
return torch.as_tensor(pickle.load(hFile)[field], dtype=torch.float).to(device)
|
|
return None
|
|
|
|
|
|
def load_mesh_auxiliary_data(
|
|
fpath: str, device: Optional[torch.device] = None
|
|
) -> Optional[torch.Tensor]:
|
|
fpath_local = PathManager.get_local_path(fpath)
|
|
with PathManager.open(fpath_local, "rb") as hFile:
|
|
return torch.as_tensor(pickle.load(hFile), dtype=torch.float).to(device)
|
|
return None
|
|
|
|
|
|
@lru_cache()
|
|
def load_mesh_symmetry(
|
|
symmetry_fpath: str, device: Optional[torch.device] = None
|
|
) -> Optional[Dict[str, torch.Tensor]]:
|
|
with PathManager.open(symmetry_fpath, "rb") as hFile:
|
|
symmetry_loaded = pickle.load(hFile)
|
|
symmetry = {
|
|
"vertex_transforms": torch.as_tensor(
|
|
symmetry_loaded["vertex_transforms"], dtype=torch.long
|
|
).to(device),
|
|
}
|
|
return symmetry
|
|
return None
|
|
|
|
|
|
@lru_cache()
|
|
def create_mesh(mesh_name: str, device: Optional[torch.device] = None) -> Mesh:
|
|
return Mesh(mesh_info=MeshCatalog[mesh_name], device=device)
|
|
|
