bluestyle97 commited on
Commit
772e07e
1 Parent(s): c778773

Delete freesplatter/utils/mesh.py

Browse files
Files changed (1) hide show
  1. 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