Spaces:
Paused
Paused
File size: 2,073 Bytes
1ac2018 |
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 |
import os
import torch
import trimesh
import numpy as np
import gc
def quad_to_triangle_mesh(F):
"""
Converts a quad-dominant mesh into a pure triangle mesh by splitting quads into two triangles.
Parameters:
quad_mesh (trimesh.Trimesh): Input mesh with quad faces.
Returns:
trimesh.Trimesh: A new mesh with only triangle faces.
"""
faces = F
### If already a triangle mesh -- skip
if len(faces[0]) == 3:
return F
new_faces = []
for face in faces:
if len(face) == 4: # Quad face
# Split into two triangles
new_faces.append([face[0], face[1], face[2]]) # Triangle 1
new_faces.append([face[0], face[2], face[3]]) # Triangle 2
else:
print(f"Warning: Skipping non-triangle/non-quad face {face}")
new_faces = np.array(new_faces)
return new_faces
class Demo_Dataset(torch.utils.data.Dataset):
def __init__(self, obj_path):
super().__init__()
self.obj_path = obj_path
self.pc_num_pts = 100000
def __len__(self):
return 1
def get_model(self):
uid = os.path.basename(self.obj_path).split(".")[-2]
mesh = trimesh.load(self.obj_path, force='mesh', process=False)
vertices = mesh.vertices
faces = mesh.faces
bbmin = vertices.min(0)
bbmax = vertices.max(0)
center = (bbmin + bbmax) * 0.5
scale = 2.0 * 0.9 / (bbmax - bbmin).max()
vertices = (vertices - center) * scale
mesh.vertices = vertices
### Make sure it is a triangle mesh -- just convert the quad
mesh.faces = quad_to_triangle_mesh(faces)
pc, _ = trimesh.sample.sample_surface(mesh, self.pc_num_pts)
result = {
'uid': uid
}
result['pc'] = torch.tensor(pc, dtype=torch.float32)
result['vertices'] = mesh.vertices
result['faces'] = mesh.faces
return result
def __getitem__(self, index):
gc.collect()
return self.get_model() |