-
Notifications
You must be signed in to change notification settings - Fork 243
Expand file tree
/
Copy pathmesh.py
More file actions
150 lines (112 loc) · 4.93 KB
/
Copy pathmesh.py
File metadata and controls
150 lines (112 loc) · 4.93 KB
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
from enum import Enum
from typing import TYPE_CHECKING, Union
import numpy as np
if TYPE_CHECKING: # pragma: no cover
from docarray.typing import T
import trimesh
class MeshEnum(Enum):
FILE_EXTENSIONS = [
'glb',
'obj',
'ply',
]
VERTICES = 'vertices'
FACES = 'faces'
class PointCloudEnum(Enum):
COLORS = 'point_cloud_colors'
class MeshDataMixin:
"""Provide helper functions for :class:`Document` to support 3D mesh data and point cloud."""
def _load_mesh(
self, force: str = None
) -> Union['trimesh.Trimesh', 'trimesh.Scene']:
"""Load a trimesh.Mesh or trimesh.Scene object from :attr:`.uri`.
:param force: str or None. For 'mesh' try to coerce scenes into a single mesh. For 'scene'
try to coerce everything into a scene.
:return: trimesh.Mesh or trimesh.Scene object
"""
import urllib.parse
import trimesh
scheme = urllib.parse.urlparse(self.uri).scheme
loader = trimesh.load_remote if scheme in ['http', 'https'] else trimesh.load
mesh = loader(self.uri, force=force)
return mesh
def load_uri_to_point_cloud_tensor(
self: 'T', samples: int, as_chunks: bool = False
) -> 'T':
"""Convert a 3d mesh-like :attr:`.uri` into :attr:`.tensor`
:param samples: number of points to sample from the mesh
:param as_chunks: when multiple geometry stored in one mesh file,
then store each geometry into different :attr:`.chunks`
:return: itself after processed
"""
if as_chunks:
import trimesh
from docarray.document import Document
# try to coerce everything into a scene
scene = self._load_mesh(force='scene')
for geo in scene.geometry.values():
geo: trimesh.Trimesh
self.chunks.append(Document(tensor=np.array(geo.sample(samples))))
else:
# combine a scene into a single mesh
mesh = self._load_mesh(force='mesh')
self.tensor = np.array(mesh.sample(samples))
return self
def load_uri_to_vertices_and_faces(self: 'T') -> 'T':
"""Convert a 3d mesh-like :attr:`.uri` into :attr:`.chunks` as vertices and faces
:return: itself after processed
"""
from docarray.document import Document
mesh = self._load_mesh(force='mesh')
vertices = mesh.vertices.view(np.ndarray)
faces = mesh.faces.view(np.ndarray)
self.chunks = [
Document(name=MeshEnum.VERTICES.value, tensor=vertices),
Document(name=MeshEnum.FACES.value, tensor=faces),
]
return self
def load_vertices_and_faces_to_point_cloud(self: 'T', samples: int) -> 'T':
"""Convert a 3d mesh of vertices and faces from :attr:`.chunks` into point cloud :attr:`.tensor`
:param samples: number of points to sample from the mesh
:return: itself after processed
"""
vertices = None
faces = None
for chunk in self.chunks:
if chunk.tags['name'] == MeshEnum.VERTICES.value:
vertices = chunk.tensor
if chunk.tags['name'] == MeshEnum.FACES.value:
faces = chunk.tensor
if vertices is not None and faces is not None:
import trimesh
mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
self.tensor = np.array(mesh.sample(samples))
else:
raise AttributeError(
'Point cloud tensor can not be set, since vertices and faces chunk tensor have not been set.'
)
return self
def load_uris_to_rgbd_tensor(self: 'T') -> 'T':
"""Load RGB image from :attr:`.uri` of :attr:`.chunks[0]` and depth image from :attr:`.uri` of :attr:`.chunks[1]` and merge them into :attr:`.tensor`.
:return: itself after processed
"""
from PIL import Image
if len(self.chunks) != 2:
raise ValueError(
f'The provided Document does not have two chunks but instead {len(self.chunks)}. To load uris to RGBD tensor, the Document needs to have two chunks, with the first one providing the RGB image uri, and the second one providing the depth image uri.'
)
for chunk in self.chunks:
if chunk.uri == '':
raise ValueError(
'A chunk of the given Document does not provide a uri.'
)
rgb_img = np.array(Image.open(self.chunks[0].uri).convert('RGB'))
depth_img = np.array(Image.open(self.chunks[1].uri))
if rgb_img.shape[0:2] != depth_img.shape:
raise ValueError(
f'The provided RGB image and depth image are not of the same shapes: {rgb_img.shape[0:2]} != {depth_img.shape}'
)
self.tensor = np.concatenate(
(rgb_img, np.expand_dims(depth_img, axis=2)), axis=-1
)
return self