| | |
| |
|
| | 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) |
| |
|
