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