Spaces:
kadirnar
/
Runtime error

File size: 6,369 Bytes
938e515
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved

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):
        # TODO: compute using Laplace-Beltrami
        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:
        # pyre-fixme[7]: Expected `Tuple[Optional[Tensor], Optional[Tensor]]` but
        #  got `Tensor`.
        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)