Spaces:
Running
on
Zero
Running
on
Zero
bluestyle97
commited on
Commit
•
772e07e
1
Parent(s):
c778773
Delete freesplatter/utils/mesh.py
Browse files- freesplatter/utils/mesh.py +0 -736
freesplatter/utils/mesh.py
DELETED
@@ -1,736 +0,0 @@
|
|
1 |
-
import os
|
2 |
-
import math
|
3 |
-
import numpy as np
|
4 |
-
import cv2
|
5 |
-
import torch
|
6 |
-
import trimesh
|
7 |
-
import torch.nn.functional as F
|
8 |
-
import pygltflib
|
9 |
-
import xatlas
|
10 |
-
import miniball
|
11 |
-
from trimesh.visual import TextureVisuals
|
12 |
-
from PIL import Image
|
13 |
-
|
14 |
-
|
15 |
-
def dot(x, y):
|
16 |
-
return torch.sum(x * y, -1, keepdim=True)
|
17 |
-
|
18 |
-
|
19 |
-
def length(x, eps=1e-20):
|
20 |
-
return torch.sqrt(torch.clamp(dot(x, x), min=eps))
|
21 |
-
|
22 |
-
|
23 |
-
def safe_normalize(x, eps=1e-20):
|
24 |
-
return x / length(x, eps)
|
25 |
-
|
26 |
-
|
27 |
-
class Mesh:
|
28 |
-
def __init__(
|
29 |
-
self,
|
30 |
-
v=None,
|
31 |
-
f=None,
|
32 |
-
vn=None,
|
33 |
-
fn=None,
|
34 |
-
vt=None,
|
35 |
-
ft=None,
|
36 |
-
vc=None,
|
37 |
-
albedo=None,
|
38 |
-
device=None,
|
39 |
-
textureless=False):
|
40 |
-
self.device = device
|
41 |
-
self.v = v
|
42 |
-
self.vn = vn
|
43 |
-
self.vt = vt
|
44 |
-
self.vc = vc
|
45 |
-
self.f = f
|
46 |
-
self.fn = fn
|
47 |
-
self.ft = ft
|
48 |
-
self.face_normals = None
|
49 |
-
# only support a single albedo
|
50 |
-
self.albedo = albedo
|
51 |
-
self.textureless = textureless
|
52 |
-
|
53 |
-
self.ori_center = 0
|
54 |
-
self.ori_scale = 1
|
55 |
-
|
56 |
-
def detach(self):
|
57 |
-
self.v = self.v.detach() if self.v is not None else None
|
58 |
-
self.vn = self.vn.detach() if self.vn is not None else None
|
59 |
-
self.vt = self.vt.detach() if self.vt is not None else None
|
60 |
-
self.vc = self.vc.detach() if self.vc is not None else None
|
61 |
-
self.f = self.f.detach() if self.f is not None else None
|
62 |
-
self.fn = self.fn.detach() if self.fn is not None else None
|
63 |
-
self.ft = self.ft.detach() if self.ft is not None else None
|
64 |
-
self.face_normals = self.face_normals.detach() if self.face_normals is not None else None
|
65 |
-
self.albedo = self.albedo.detach() if self.albedo is not None else None
|
66 |
-
return self
|
67 |
-
|
68 |
-
@classmethod
|
69 |
-
def load(cls, path=None, resize=False, auto_uv=True, flip_yz=False, **kwargs):
|
70 |
-
# assume init with kwargs
|
71 |
-
if path is None:
|
72 |
-
mesh = cls(**kwargs)
|
73 |
-
# obj supports face uv
|
74 |
-
elif path.endswith(".obj"):
|
75 |
-
mesh = cls.load_obj(path, **kwargs)
|
76 |
-
# trimesh only supports vertex uv, but can load more formats
|
77 |
-
else:
|
78 |
-
mesh = cls.load_trimesh(path, **kwargs)
|
79 |
-
|
80 |
-
print(f"[Mesh loading] v: {mesh.v.shape}, f: {mesh.f.shape}")
|
81 |
-
# auto-normalize
|
82 |
-
if resize:
|
83 |
-
mesh.auto_size()
|
84 |
-
# auto-fix normal
|
85 |
-
if mesh.vn is None:
|
86 |
-
mesh.auto_normal()
|
87 |
-
print(f"[Mesh loading] vn: {mesh.vn.shape}, fn: {mesh.fn.shape}")
|
88 |
-
# auto-fix texture
|
89 |
-
if mesh.vt is None and auto_uv:
|
90 |
-
mesh.auto_uv(cache_path=path)
|
91 |
-
if mesh.vt is not None and mesh.ft is not None:
|
92 |
-
print(f"[Mesh loading] vt: {mesh.vt.shape}, ft: {mesh.ft.shape}")
|
93 |
-
|
94 |
-
if flip_yz:
|
95 |
-
mesh.v[..., [1, 2]] = mesh.v[..., [2, 1]]
|
96 |
-
mesh.vn[..., [1, 2]] = mesh.vn[..., [2, 1]]
|
97 |
-
mesh.v[..., 1] = -mesh.v[..., 1]
|
98 |
-
mesh.vn[..., 1] = -mesh.vn[..., 1]
|
99 |
-
|
100 |
-
return mesh
|
101 |
-
|
102 |
-
# load from obj file
|
103 |
-
@classmethod
|
104 |
-
def load_obj(cls, path, albedo_path=None, device=None):
|
105 |
-
assert os.path.splitext(path)[-1] == ".obj"
|
106 |
-
|
107 |
-
mesh = cls()
|
108 |
-
|
109 |
-
# device
|
110 |
-
if device is None:
|
111 |
-
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
112 |
-
|
113 |
-
mesh.device = device
|
114 |
-
|
115 |
-
# load obj
|
116 |
-
with open(path, "r") as f:
|
117 |
-
lines = f.readlines()
|
118 |
-
|
119 |
-
def parse_f_v(fv):
|
120 |
-
# pass in a vertex term of a face, return {v, vt, vn} (-1 if not provided)
|
121 |
-
# supported forms:
|
122 |
-
# f v1 v2 v3
|
123 |
-
# f v1/vt1 v2/vt2 v3/vt3
|
124 |
-
# f v1/vt1/vn1 v2/vt2/vn2 v3/vt3/vn3
|
125 |
-
# f v1//vn1 v2//vn2 v3//vn3
|
126 |
-
xs = [int(x) - 1 if x != "" else -1 for x in fv.split("/")]
|
127 |
-
xs.extend([-1] * (3 - len(xs)))
|
128 |
-
return xs[0], xs[1], xs[2]
|
129 |
-
|
130 |
-
# NOTE: we ignore usemtl, and assume the mesh ONLY uses one material (first in mtl)
|
131 |
-
vertices, texcoords, normals = [], [], []
|
132 |
-
faces, tfaces, nfaces = [], [], []
|
133 |
-
mtl_path = None
|
134 |
-
|
135 |
-
for line in lines:
|
136 |
-
split_line = line.split()
|
137 |
-
# empty line
|
138 |
-
if len(split_line) == 0:
|
139 |
-
continue
|
140 |
-
prefix = split_line[0].lower()
|
141 |
-
# mtllib
|
142 |
-
if prefix == "mtllib":
|
143 |
-
mtl_path = split_line[1]
|
144 |
-
# usemtl
|
145 |
-
elif prefix == "usemtl":
|
146 |
-
pass # ignored
|
147 |
-
# v/vn/vt
|
148 |
-
elif prefix == "v":
|
149 |
-
vertices.append([float(v) for v in split_line[1:]])
|
150 |
-
elif prefix == "vn":
|
151 |
-
normals.append([float(v) for v in split_line[1:]])
|
152 |
-
elif prefix == "vt":
|
153 |
-
val = [float(v) for v in split_line[1:]]
|
154 |
-
texcoords.append([val[0], 1.0 - val[1]])
|
155 |
-
elif prefix == "f":
|
156 |
-
vs = split_line[1:]
|
157 |
-
nv = len(vs)
|
158 |
-
v0, t0, n0 = parse_f_v(vs[0])
|
159 |
-
for i in range(nv - 2): # triangulate (assume vertices are ordered)
|
160 |
-
v1, t1, n1 = parse_f_v(vs[i + 1])
|
161 |
-
v2, t2, n2 = parse_f_v(vs[i + 2])
|
162 |
-
faces.append([v0, v1, v2])
|
163 |
-
tfaces.append([t0, t1, t2])
|
164 |
-
nfaces.append([n0, n1, n2])
|
165 |
-
|
166 |
-
mesh.v = torch.tensor(vertices, dtype=torch.float32, device=device)
|
167 |
-
mesh.vt = (
|
168 |
-
torch.tensor(texcoords, dtype=torch.float32, device=device)
|
169 |
-
if len(texcoords) > 0
|
170 |
-
else None
|
171 |
-
)
|
172 |
-
mesh.vn = (
|
173 |
-
torch.tensor(normals, dtype=torch.float32, device=device)
|
174 |
-
if len(normals) > 0
|
175 |
-
else None
|
176 |
-
)
|
177 |
-
|
178 |
-
mesh.f = torch.tensor(faces, dtype=torch.int32, device=device)
|
179 |
-
mesh.ft = (
|
180 |
-
torch.tensor(tfaces, dtype=torch.int32, device=device)
|
181 |
-
if len(texcoords) > 0
|
182 |
-
else None
|
183 |
-
)
|
184 |
-
mesh.fn = (
|
185 |
-
torch.tensor(nfaces, dtype=torch.int32, device=device)
|
186 |
-
if len(normals) > 0
|
187 |
-
else None
|
188 |
-
)
|
189 |
-
|
190 |
-
# see if there is vertex color
|
191 |
-
if mesh.v.size(-1) > 3:
|
192 |
-
mesh.vc = mesh.v[:, 3:]
|
193 |
-
mesh.v = mesh.v[:, :3]
|
194 |
-
if mesh.vc.size(-1) == 3:
|
195 |
-
mesh.vc = torch.cat([mesh.vc, torch.ones_like(mesh.vc[:, :1])], dim=-1)
|
196 |
-
print(f"[load_obj] use vertex color: {mesh.vc.shape}")
|
197 |
-
|
198 |
-
# try to retrieve mtl file
|
199 |
-
mtl_path_candidates = []
|
200 |
-
if mtl_path is not None:
|
201 |
-
mtl_path_candidates.append(mtl_path)
|
202 |
-
mtl_path_candidates.append(os.path.join(os.path.dirname(path), mtl_path))
|
203 |
-
mtl_path_candidates.append(path.replace(".obj", ".mtl"))
|
204 |
-
|
205 |
-
mtl_path = None
|
206 |
-
for candidate in mtl_path_candidates:
|
207 |
-
if os.path.exists(candidate):
|
208 |
-
mtl_path = candidate
|
209 |
-
break
|
210 |
-
|
211 |
-
# if albedo_path is not provided, try retrieve it from mtl
|
212 |
-
if mtl_path is not None and albedo_path is None:
|
213 |
-
with open(mtl_path, "r") as f:
|
214 |
-
lines = f.readlines()
|
215 |
-
for line in lines:
|
216 |
-
split_line = line.split()
|
217 |
-
# empty line
|
218 |
-
if len(split_line) == 0:
|
219 |
-
continue
|
220 |
-
prefix = split_line[0]
|
221 |
-
# NOTE: simply use the first map_Kd as albedo!
|
222 |
-
if "map_Kd" in prefix:
|
223 |
-
albedo_path = os.path.join(os.path.dirname(path), split_line[1])
|
224 |
-
print(f"[load_obj] use texture from: {albedo_path}")
|
225 |
-
break
|
226 |
-
|
227 |
-
# still not found albedo_path, or the path doesn't exist
|
228 |
-
if albedo_path is None or not os.path.exists(albedo_path):
|
229 |
-
# init an empty texture
|
230 |
-
print(f"[load_obj] init empty albedo!")
|
231 |
-
# albedo = np.random.rand(1024, 1024, 3).astype(np.float32)
|
232 |
-
albedo = np.ones((1024, 1024, 3), dtype=np.float32) * np.array([0.5, 0.5, 0.5]) # default color
|
233 |
-
mesh.textureless = True
|
234 |
-
else:
|
235 |
-
albedo = cv2.imread(albedo_path, cv2.IMREAD_UNCHANGED)
|
236 |
-
albedo = cv2.cvtColor(albedo, cv2.COLOR_BGR2RGB)
|
237 |
-
albedo = albedo.astype(np.float32) / 255
|
238 |
-
print(f"[load_obj] load texture: {albedo.shape}")
|
239 |
-
|
240 |
-
# import matplotlib.pyplot as plt
|
241 |
-
# plt.imshow(albedo)
|
242 |
-
# plt.show()
|
243 |
-
|
244 |
-
mesh.albedo = torch.tensor(albedo, dtype=torch.float32, device=device)
|
245 |
-
|
246 |
-
return mesh
|
247 |
-
|
248 |
-
@classmethod
|
249 |
-
def load_trimesh(cls, path, device=None):
|
250 |
-
mesh = cls()
|
251 |
-
|
252 |
-
# device
|
253 |
-
if device is None:
|
254 |
-
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
255 |
-
|
256 |
-
mesh.device = device
|
257 |
-
|
258 |
-
# use trimesh to load glb, assume only has one single RootMesh...
|
259 |
-
_data = trimesh.load(path)
|
260 |
-
if isinstance(_data, trimesh.Scene):
|
261 |
-
mesh_keys = list(_data.geometry.keys())
|
262 |
-
assert (
|
263 |
-
len(mesh_keys) == 1
|
264 |
-
), f"{path} contains more than one meshes, not supported!"
|
265 |
-
_mesh = _data.geometry[mesh_keys[0]]
|
266 |
-
|
267 |
-
elif isinstance(_data, trimesh.Trimesh):
|
268 |
-
_mesh = _data
|
269 |
-
|
270 |
-
else:
|
271 |
-
raise NotImplementedError(f"type {type(_data)} not supported!")
|
272 |
-
|
273 |
-
if hasattr(_mesh.visual, "material"):
|
274 |
-
_material = _mesh.visual.material
|
275 |
-
if isinstance(_material, trimesh.visual.material.PBRMaterial):
|
276 |
-
texture = np.array(_material.baseColorTexture).astype(np.float32) / 255
|
277 |
-
elif isinstance(_material, trimesh.visual.material.SimpleMaterial):
|
278 |
-
texture = (
|
279 |
-
np.array(_material.to_pbr().baseColorTexture).astype(np.float32) / 255
|
280 |
-
)
|
281 |
-
else:
|
282 |
-
raise NotImplementedError(f"material type {type(_material)} not supported!")
|
283 |
-
|
284 |
-
print(f"[load_obj] load texture: {texture.shape}")
|
285 |
-
mesh.albedo = torch.tensor(texture, dtype=torch.float32, device=device)
|
286 |
-
|
287 |
-
if hasattr(_mesh.visual, "uv"):
|
288 |
-
texcoords = _mesh.visual.uv
|
289 |
-
texcoords[:, 1] = 1 - texcoords[:, 1]
|
290 |
-
mesh.vt = (
|
291 |
-
torch.tensor(texcoords, dtype=torch.float32, device=device)
|
292 |
-
if len(texcoords) > 0
|
293 |
-
else None
|
294 |
-
)
|
295 |
-
else:
|
296 |
-
texcoords = None
|
297 |
-
|
298 |
-
if hasattr(_mesh.visual, "vertex_colors"):
|
299 |
-
colors = _mesh.visual.vertex_colors
|
300 |
-
mesh.vc = (
|
301 |
-
torch.tensor(colors, dtype=torch.float32, device=device) / 255
|
302 |
-
if len(colors) > 0
|
303 |
-
else None
|
304 |
-
)
|
305 |
-
|
306 |
-
vertices = _mesh.vertices
|
307 |
-
|
308 |
-
normals = _mesh.vertex_normals
|
309 |
-
|
310 |
-
# trimesh only support vertex uv...
|
311 |
-
faces = tfaces = nfaces = _mesh.faces
|
312 |
-
|
313 |
-
mesh.v = torch.tensor(vertices, dtype=torch.float32, device=device)
|
314 |
-
mesh.vn = (
|
315 |
-
torch.tensor(normals, dtype=torch.float32, device=device)
|
316 |
-
if len(normals) > 0
|
317 |
-
else None
|
318 |
-
)
|
319 |
-
|
320 |
-
mesh.f = torch.tensor(faces, dtype=torch.int32, device=device)
|
321 |
-
mesh.ft = (
|
322 |
-
torch.tensor(tfaces, dtype=torch.int32, device=device)
|
323 |
-
if texcoords is not None
|
324 |
-
else None
|
325 |
-
)
|
326 |
-
mesh.fn = (
|
327 |
-
torch.tensor(nfaces, dtype=torch.int32, device=device)
|
328 |
-
if normals is not None
|
329 |
-
else None
|
330 |
-
)
|
331 |
-
|
332 |
-
return mesh
|
333 |
-
|
334 |
-
# aabb
|
335 |
-
def aabb(self):
|
336 |
-
return torch.min(self.v, dim=0).values, torch.max(self.v, dim=0).values
|
337 |
-
|
338 |
-
# unit size
|
339 |
-
@torch.no_grad()
|
340 |
-
def auto_size(self):
|
341 |
-
vmin, vmax = self.aabb()
|
342 |
-
self.ori_center = (vmax + vmin) / 2
|
343 |
-
self.ori_scale = 1.2 / torch.max(vmax - vmin).item() # to ~ [-0.6, 0.6]
|
344 |
-
self.v = (self.v - self.ori_center) * self.ori_scale
|
345 |
-
|
346 |
-
def auto_normal(self):
|
347 |
-
i0, i1, i2 = self.f[:, 0].long(), self.f[:, 1].long(), self.f[:, 2].long()
|
348 |
-
v0, v1, v2 = self.v[i0, :], self.v[i1, :], self.v[i2, :]
|
349 |
-
|
350 |
-
face_normals = torch.cross(v1 - v0, v2 - v0)
|
351 |
-
|
352 |
-
# Splat face normals to vertices
|
353 |
-
face_normals = F.normalize(face_normals, dim=-1)
|
354 |
-
vn = torch.zeros_like(self.v)
|
355 |
-
vn.scatter_add_(0, i0[:, None].repeat(1, 3), face_normals)
|
356 |
-
vn.scatter_add_(0, i1[:, None].repeat(1, 3), face_normals)
|
357 |
-
vn.scatter_add_(0, i2[:, None].repeat(1, 3), face_normals)
|
358 |
-
vn = F.normalize(vn, dim=-1)
|
359 |
-
|
360 |
-
self.vn = vn
|
361 |
-
self.fn = self.f
|
362 |
-
self.face_normals = face_normals
|
363 |
-
|
364 |
-
def auto_uv(self, cache_path=None, vmap=True):
|
365 |
-
# try to load cache
|
366 |
-
if cache_path is not None:
|
367 |
-
cache_path = os.path.splitext(cache_path)[0] + '_uv.npz'
|
368 |
-
|
369 |
-
if cache_path is not None and os.path.exists(cache_path):
|
370 |
-
data = np.load(cache_path)
|
371 |
-
vt_np, ft_np, vmapping = data['vt'], data['ft'], data['vmapping']
|
372 |
-
else:
|
373 |
-
v_np = self.v.detach().cpu().numpy()
|
374 |
-
f_np = self.f.detach().int().cpu().numpy()
|
375 |
-
atlas = xatlas.Atlas()
|
376 |
-
atlas.add_mesh(v_np, f_np)
|
377 |
-
chart_options = xatlas.ChartOptions()
|
378 |
-
# chart_options.max_iterations = 4
|
379 |
-
atlas.generate(chart_options=chart_options)
|
380 |
-
vmapping, ft_np, vt_np = atlas[0] # [N], [M, 3], [N, 2]
|
381 |
-
|
382 |
-
# save to cache
|
383 |
-
if cache_path is not None:
|
384 |
-
np.savez(cache_path, vt=vt_np, ft=ft_np, vmapping=vmapping)
|
385 |
-
|
386 |
-
vt = torch.from_numpy(vt_np.astype(np.float32)).to(self.device)
|
387 |
-
ft = torch.from_numpy(ft_np.astype(np.int32)).to(self.device)
|
388 |
-
self.vt = vt
|
389 |
-
self.ft = ft
|
390 |
-
|
391 |
-
if vmap:
|
392 |
-
# remap v/f to vt/ft, so each v correspond to a unique vt. (necessary for gltf)
|
393 |
-
vmapping = torch.from_numpy(vmapping.astype(np.int64)).long().to(self.device)
|
394 |
-
self.align_v_to_vt(vmapping)
|
395 |
-
|
396 |
-
def align_v_to_vt(self, vmapping=None):
|
397 |
-
# remap v/f and vn/vn to vt/ft.
|
398 |
-
if vmapping is None:
|
399 |
-
ft = self.ft.view(-1).long()
|
400 |
-
f = self.f.view(-1).long()
|
401 |
-
vmapping = torch.zeros(self.vt.shape[0], dtype=torch.long, device=self.device)
|
402 |
-
vmapping[ft] = f # scatter, randomly choose one if index is not unique
|
403 |
-
if self.vn is not None and (self.f == self.fn).all():
|
404 |
-
self.vn = self.vn[vmapping]
|
405 |
-
self.fn = self.ft
|
406 |
-
self.v = self.v[vmapping]
|
407 |
-
self.f = self.ft
|
408 |
-
|
409 |
-
def align_vn_to_vt(self, vmapping=None):
|
410 |
-
if vmapping is None:
|
411 |
-
ft = self.ft.view(-1).long()
|
412 |
-
fn = self.f.view(-1).long()
|
413 |
-
vmapping = torch.zeros(self.vt.shape[0], dtype=torch.long, device=self.device)
|
414 |
-
vmapping[ft] = fn # scatter, randomly choose one if index is not unique
|
415 |
-
self.vn = self.vn[vmapping]
|
416 |
-
self.fn = self.ft
|
417 |
-
|
418 |
-
def to(self, device):
|
419 |
-
self.device = device
|
420 |
-
for name in ['v', 'f', 'vn', 'fn', 'vt', 'ft', 'albedo', 'vc', 'face_normals']:
|
421 |
-
tensor = getattr(self, name)
|
422 |
-
if tensor is not None:
|
423 |
-
setattr(self, name, tensor.to(device))
|
424 |
-
return self
|
425 |
-
|
426 |
-
def copy(self):
|
427 |
-
return Mesh(
|
428 |
-
v=self.v,
|
429 |
-
f=self.f,
|
430 |
-
vn=self.vn,
|
431 |
-
fn=self.fn,
|
432 |
-
vt=self.vt,
|
433 |
-
ft=self.ft,
|
434 |
-
vc=self.vc,
|
435 |
-
albedo=self.albedo,
|
436 |
-
device=self.device,
|
437 |
-
textureless=self.textureless)
|
438 |
-
|
439 |
-
def write(self, path, flip_yz=False):
|
440 |
-
mesh = self.copy()
|
441 |
-
if flip_yz:
|
442 |
-
mesh.v = mesh.v.clone()
|
443 |
-
mesh.vn = mesh.vn.clone()
|
444 |
-
mesh.v[..., 1] = -mesh.v[..., 1]
|
445 |
-
mesh.vn[..., 1] = -mesh.vn[..., 1]
|
446 |
-
mesh.v[..., [1, 2]] = mesh.v[..., [2, 1]]
|
447 |
-
mesh.vn[..., [1, 2]] = mesh.vn[..., [2, 1]]
|
448 |
-
if path.endswith('.ply'):
|
449 |
-
mesh.write_ply(path)
|
450 |
-
elif path.endswith('.obj'):
|
451 |
-
mesh.write_obj(path)
|
452 |
-
elif path.endswith('.glb') or path.endswith('.gltf'):
|
453 |
-
mesh.write_glb(path)
|
454 |
-
else:
|
455 |
-
raise NotImplementedError(f'format {path} not supported!')
|
456 |
-
|
457 |
-
# write to ply file (only geom)
|
458 |
-
def write_ply(self, path):
|
459 |
-
|
460 |
-
v_np = self.v.detach().cpu().numpy()
|
461 |
-
f_np = self.f.detach().cpu().numpy()
|
462 |
-
|
463 |
-
_mesh = trimesh.Trimesh(vertices=v_np, faces=f_np)
|
464 |
-
_mesh.export(path)
|
465 |
-
|
466 |
-
# write to gltf/glb file (geom + texture)
|
467 |
-
def write_glb(self, path):
|
468 |
-
|
469 |
-
assert self.vn is not None
|
470 |
-
if self.vt is None:
|
471 |
-
self.vt = self.v.new_zeros((self.v.size(0), 2))
|
472 |
-
self.ft = self.f
|
473 |
-
if (self.f != self.ft).any():
|
474 |
-
self.align_v_to_vt()
|
475 |
-
if (self.fn != self.ft).any():
|
476 |
-
self.align_vn_to_vt()
|
477 |
-
|
478 |
-
assert self.v.shape[0] == self.vn.shape[0] and self.v.shape[0] == self.vt.shape[0]
|
479 |
-
|
480 |
-
f_np = self.f.detach().cpu().numpy().astype(np.uint32)
|
481 |
-
v_np = self.v.detach().cpu().numpy().astype(np.float32)
|
482 |
-
vt_np = self.vt.detach().cpu().numpy().astype(np.float32)
|
483 |
-
vn_np = self.vn.detach().cpu().numpy().astype(np.float32)
|
484 |
-
|
485 |
-
albedo = self.albedo.detach().cpu().numpy() if self.albedo is not None \
|
486 |
-
else np.full((1024, 1024, 3), 0.5, dtype=np.float32)
|
487 |
-
albedo = (albedo * 255).astype(np.uint8)
|
488 |
-
albedo = cv2.cvtColor(albedo, cv2.COLOR_RGB2BGR)
|
489 |
-
|
490 |
-
f_np_blob = f_np.flatten().tobytes()
|
491 |
-
v_np_blob = v_np.tobytes()
|
492 |
-
vt_np_blob = vt_np.tobytes()
|
493 |
-
vn_np_blob = vn_np.tobytes()
|
494 |
-
albedo_blob = cv2.imencode('.png', albedo)[1].tobytes()
|
495 |
-
|
496 |
-
gltf = pygltflib.GLTF2(
|
497 |
-
scene=0,
|
498 |
-
scenes=[pygltflib.Scene(nodes=[0])],
|
499 |
-
nodes=[pygltflib.Node(mesh=0)],
|
500 |
-
meshes=[pygltflib.Mesh(primitives=[
|
501 |
-
pygltflib.Primitive(
|
502 |
-
# indices to accessors (0 is triangles)
|
503 |
-
attributes=pygltflib.Attributes(
|
504 |
-
POSITION=1, TEXCOORD_0=2, NORMAL=3
|
505 |
-
),
|
506 |
-
indices=0, material=0,
|
507 |
-
)
|
508 |
-
])],
|
509 |
-
materials=[
|
510 |
-
pygltflib.Material(
|
511 |
-
pbrMetallicRoughness=pygltflib.PbrMetallicRoughness(
|
512 |
-
baseColorTexture=pygltflib.TextureInfo(index=0, texCoord=0),
|
513 |
-
metallicFactor=0.0,
|
514 |
-
roughnessFactor=1.0,
|
515 |
-
),
|
516 |
-
alphaCutoff=0,
|
517 |
-
doubleSided=True,
|
518 |
-
)
|
519 |
-
],
|
520 |
-
textures=[
|
521 |
-
pygltflib.Texture(sampler=0, source=0),
|
522 |
-
],
|
523 |
-
samplers=[
|
524 |
-
pygltflib.Sampler(magFilter=pygltflib.LINEAR, minFilter=pygltflib.LINEAR_MIPMAP_LINEAR,
|
525 |
-
wrapS=pygltflib.REPEAT, wrapT=pygltflib.REPEAT),
|
526 |
-
],
|
527 |
-
images=[
|
528 |
-
# use embedded (buffer) image
|
529 |
-
pygltflib.Image(bufferView=4, mimeType="image/png"),
|
530 |
-
],
|
531 |
-
buffers=[
|
532 |
-
pygltflib.Buffer(
|
533 |
-
byteLength=len(f_np_blob) + len(v_np_blob) + len(vt_np_blob) + len(vn_np_blob) + len(albedo_blob))
|
534 |
-
],
|
535 |
-
# buffer view (based on dtype)
|
536 |
-
bufferViews=[
|
537 |
-
# triangles; as flatten (element) array
|
538 |
-
pygltflib.BufferView(
|
539 |
-
buffer=0,
|
540 |
-
byteLength=len(f_np_blob),
|
541 |
-
target=pygltflib.ELEMENT_ARRAY_BUFFER, # GL_ELEMENT_ARRAY_BUFFER (34963)
|
542 |
-
),
|
543 |
-
# positions; as vec3 array
|
544 |
-
pygltflib.BufferView(
|
545 |
-
buffer=0,
|
546 |
-
byteOffset=len(f_np_blob),
|
547 |
-
byteLength=len(v_np_blob),
|
548 |
-
byteStride=12, # vec3
|
549 |
-
target=pygltflib.ARRAY_BUFFER, # GL_ARRAY_BUFFER (34962)
|
550 |
-
),
|
551 |
-
# texcoords; as vec2 array
|
552 |
-
pygltflib.BufferView(
|
553 |
-
buffer=0,
|
554 |
-
byteOffset=len(f_np_blob) + len(v_np_blob),
|
555 |
-
byteLength=len(vt_np_blob),
|
556 |
-
byteStride=8, # vec2
|
557 |
-
target=pygltflib.ARRAY_BUFFER,
|
558 |
-
),
|
559 |
-
# normals; as vec3 array
|
560 |
-
pygltflib.BufferView(
|
561 |
-
buffer=0,
|
562 |
-
byteOffset=len(f_np_blob) + len(v_np_blob) + len(vt_np_blob),
|
563 |
-
byteLength=len(vn_np_blob),
|
564 |
-
byteStride=12, # vec3
|
565 |
-
target=pygltflib.ARRAY_BUFFER,
|
566 |
-
),
|
567 |
-
# texture; as none target
|
568 |
-
pygltflib.BufferView(
|
569 |
-
buffer=0,
|
570 |
-
byteOffset=len(f_np_blob) + len(v_np_blob) + len(vt_np_blob) + len(vn_np_blob),
|
571 |
-
byteLength=len(albedo_blob),
|
572 |
-
),
|
573 |
-
],
|
574 |
-
accessors=[
|
575 |
-
# 0 = triangles
|
576 |
-
pygltflib.Accessor(
|
577 |
-
bufferView=0,
|
578 |
-
componentType=pygltflib.UNSIGNED_INT, # GL_UNSIGNED_INT (5125)
|
579 |
-
count=f_np.size,
|
580 |
-
type=pygltflib.SCALAR,
|
581 |
-
max=[int(f_np.max())],
|
582 |
-
min=[int(f_np.min())],
|
583 |
-
),
|
584 |
-
# 1 = positions
|
585 |
-
pygltflib.Accessor(
|
586 |
-
bufferView=1,
|
587 |
-
componentType=pygltflib.FLOAT, # GL_FLOAT (5126)
|
588 |
-
count=len(v_np),
|
589 |
-
type=pygltflib.VEC3,
|
590 |
-
max=v_np.max(axis=0).tolist(),
|
591 |
-
min=v_np.min(axis=0).tolist(),
|
592 |
-
),
|
593 |
-
# 2 = texcoords
|
594 |
-
pygltflib.Accessor(
|
595 |
-
bufferView=2,
|
596 |
-
componentType=pygltflib.FLOAT,
|
597 |
-
count=len(vt_np),
|
598 |
-
type=pygltflib.VEC2,
|
599 |
-
max=vt_np.max(axis=0).tolist(),
|
600 |
-
min=vt_np.min(axis=0).tolist(),
|
601 |
-
),
|
602 |
-
# 3 = normals
|
603 |
-
pygltflib.Accessor(
|
604 |
-
bufferView=3,
|
605 |
-
componentType=pygltflib.FLOAT,
|
606 |
-
count=len(vn_np),
|
607 |
-
type=pygltflib.VEC3,
|
608 |
-
max=vn_np.max(axis=0).tolist(),
|
609 |
-
min=vn_np.min(axis=0).tolist(),
|
610 |
-
),
|
611 |
-
],
|
612 |
-
)
|
613 |
-
|
614 |
-
# set actual data
|
615 |
-
gltf.set_binary_blob(f_np_blob + v_np_blob + vt_np_blob + vn_np_blob + albedo_blob)
|
616 |
-
|
617 |
-
# glb = b"".join(gltf.save_to_bytes())
|
618 |
-
gltf.save(path)
|
619 |
-
|
620 |
-
# write to obj file (geom + texture)
|
621 |
-
def write_obj(self, path):
|
622 |
-
|
623 |
-
mtl_path = path.replace(".obj", ".mtl")
|
624 |
-
albedo_path = path.replace(".obj", "_albedo.png")
|
625 |
-
|
626 |
-
v_np = self.v.detach().cpu().numpy()
|
627 |
-
vt_np = self.vt.detach().cpu().numpy() if self.vt is not None else None
|
628 |
-
vn_np = self.vn.detach().cpu().numpy() if self.vn is not None else None
|
629 |
-
f_np = self.f.detach().cpu().numpy()
|
630 |
-
ft_np = self.ft.detach().cpu().numpy() if self.ft is not None else None
|
631 |
-
fn_np = self.fn.detach().cpu().numpy() if self.fn is not None else None
|
632 |
-
|
633 |
-
with open(path, "w") as fp:
|
634 |
-
fp.write(f"mtllib {os.path.basename(mtl_path)} \n")
|
635 |
-
|
636 |
-
for v in v_np:
|
637 |
-
fp.write(f"v {v[0]:.6f} {v[1]:.6f} {v[2]:.6f} \n")
|
638 |
-
|
639 |
-
if vt_np is not None:
|
640 |
-
for v in vt_np:
|
641 |
-
fp.write(f"vt {v[0]:.4f} {1 - v[1]:.4f} \n")
|
642 |
-
|
643 |
-
if vn_np is not None:
|
644 |
-
for v in vn_np:
|
645 |
-
fp.write(f"vn {v[0]:.4f} {v[1]:.4f} {v[2]:.4f} \n")
|
646 |
-
|
647 |
-
fp.write(f"usemtl defaultMat \n")
|
648 |
-
for i in range(len(f_np)):
|
649 |
-
fp.write(
|
650 |
-
f'f {f_np[i, 0] + 1}/{ft_np[i, 0] + 1 if ft_np is not None else ""}/{fn_np[i, 0] + 1 if fn_np is not None else ""} \
|
651 |
-
{f_np[i, 1] + 1}/{ft_np[i, 1] + 1 if ft_np is not None else ""}/{fn_np[i, 1] + 1 if fn_np is not None else ""} \
|
652 |
-
{f_np[i, 2] + 1}/{ft_np[i, 2] + 1 if ft_np is not None else ""}/{fn_np[i, 2] + 1 if fn_np is not None else ""} \n'
|
653 |
-
)
|
654 |
-
|
655 |
-
with open(mtl_path, "w") as fp:
|
656 |
-
fp.write(f"newmtl defaultMat \n")
|
657 |
-
fp.write(f"Ka 1 1 1 \n")
|
658 |
-
fp.write(f"Kd 1 1 1 \n")
|
659 |
-
fp.write(f"Ks 0 0 0 \n")
|
660 |
-
fp.write(f"Tr 1 \n")
|
661 |
-
fp.write(f"illum 1 \n")
|
662 |
-
fp.write(f"Ns 0 \n")
|
663 |
-
if not self.textureless and self.albedo is not None:
|
664 |
-
fp.write(f"map_Kd {os.path.basename(albedo_path)} \n")
|
665 |
-
|
666 |
-
if not self.textureless and self.albedo is not None:
|
667 |
-
albedo = self.albedo.detach().cpu().numpy()
|
668 |
-
albedo = (albedo * 255).astype(np.uint8)
|
669 |
-
cv2.imwrite(albedo_path, cv2.cvtColor(albedo, cv2.COLOR_RGB2BGR))
|
670 |
-
|
671 |
-
|
672 |
-
def normalize_mesh(mesh, tgt_radius=0.9):
|
673 |
-
mb = miniball.Miniball(mesh.v[:, :2].cpu().numpy())
|
674 |
-
center_xy = mb.center()
|
675 |
-
radius_xy_sq = mb.squared_radius()
|
676 |
-
max_z = mesh.v[:, 2].max().item()
|
677 |
-
min_z = mesh.v[:, 2].min().item()
|
678 |
-
center = mesh.v.new_tensor([center_xy[0], center_xy[1], (max_z + min_z) / 2])
|
679 |
-
radius = max(math.sqrt(radius_xy_sq), (max_z - min_z) / 2)
|
680 |
-
scale = tgt_radius / radius
|
681 |
-
mesh.v = (mesh.v - center) * scale
|
682 |
-
return mesh, center.tolist(), scale
|
683 |
-
|
684 |
-
|
685 |
-
def check_has_texture_single(geom):
|
686 |
-
return isinstance(geom.visual, TextureVisuals) and geom.visual.material.baseColorTexture is not None
|
687 |
-
|
688 |
-
|
689 |
-
def check_has_texture(mesh):
|
690 |
-
if isinstance(mesh, trimesh.Scene):
|
691 |
-
has_texture = []
|
692 |
-
for geom in mesh.geometry.values():
|
693 |
-
has_texture.append(check_has_texture_single(geom))
|
694 |
-
elif isinstance(mesh, trimesh.Trimesh):
|
695 |
-
has_texture = check_has_texture_single(mesh)
|
696 |
-
else:
|
697 |
-
raise NotImplementedError(f"type {type(mesh)} not supported!")
|
698 |
-
return has_texture
|
699 |
-
|
700 |
-
|
701 |
-
def create_texture(geom):
|
702 |
-
if hasattr(geom.visual, 'material') and hasattr(geom.visual.material, 'main_color'):
|
703 |
-
main_color = tuple(geom.visual.material.main_color)
|
704 |
-
else:
|
705 |
-
main_color = (128, 128, 128)
|
706 |
-
geom.visual = trimesh.visual.TextureVisuals(
|
707 |
-
uv=np.full((geom.vertices.shape[0], 2), 0.5),
|
708 |
-
material=trimesh.visual.material.PBRMaterial(
|
709 |
-
baseColorTexture=Image.new('RGB', (8, 8), main_color)
|
710 |
-
)
|
711 |
-
)
|
712 |
-
|
713 |
-
|
714 |
-
def color_to_texture(mesh):
|
715 |
-
if isinstance(mesh, trimesh.Scene):
|
716 |
-
for geom in mesh.geometry.values():
|
717 |
-
if not check_has_texture_single(geom):
|
718 |
-
create_texture(geom)
|
719 |
-
elif isinstance(mesh, trimesh.Trimesh):
|
720 |
-
if not check_has_texture_single(mesh):
|
721 |
-
create_texture(mesh)
|
722 |
-
else:
|
723 |
-
raise NotImplementedError(f"type {type(mesh)} not supported!")
|
724 |
-
return mesh
|
725 |
-
|
726 |
-
|
727 |
-
def purge_scene(scene):
|
728 |
-
update_flag = False
|
729 |
-
delete_list = []
|
730 |
-
for name, geom in scene.geometry.items():
|
731 |
-
if not isinstance(geom, trimesh.Trimesh):
|
732 |
-
update_flag = True
|
733 |
-
delete_list.append(name)
|
734 |
-
for name in delete_list:
|
735 |
-
scene.delete_geometry(name)
|
736 |
-
return update_flag
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|