-
Notifications
You must be signed in to change notification settings - Fork 11
/
Copy pathmesh.py
113 lines (94 loc) · 4.72 KB
/
mesh.py
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
import kaolin as kal
import torch
import utils
from utils import device
import copy
import numpy as np
import PIL
class Mesh():
def __init__(self,obj_path,color=torch.tensor([0.0,0.0,1.0])):
if ".obj" in obj_path:
mesh = kal.io.obj.import_mesh(obj_path, with_normals=True)
elif ".off" in obj_path:
mesh = kal.io.off.import_mesh(obj_path)
else:
raise ValueError(f"{obj_path} extension not implemented in mesh reader.")
self.vertices = mesh.vertices.to(device)
self.faces = mesh.faces.to(device)
self.vertex_normals = None
self.face_normals = None
self.texture_map = None
self.face_uvs = None
if ".obj" in obj_path:
if mesh.vertex_normals is not None:
self.vertex_normals = mesh.vertex_normals.to(device).float()
# Normalize
self.vertex_normals = torch.nn.functional.normalize(self.vertex_normals)
if mesh.face_normals is not None:
self.face_normals = mesh.face_normals.to(device).float()
# Normalize
self.face_normals = torch.nn.functional.normalize(self.face_normals)
self.set_mesh_color(color)
def standardize_mesh(self,inplace=False):
mesh = self if inplace else copy.deepcopy(self)
return utils.standardize_mesh(mesh)
def normalize_mesh(self,inplace=False):
mesh = self if inplace else copy.deepcopy(self)
return utils.normalize_mesh(mesh)
def update_vertex(self,verts,inplace=False):
mesh = self if inplace else copy.deepcopy(self)
mesh.vertices = verts
return mesh
def set_mesh_color(self,color):
self.texture_map = utils.get_texture_map_from_color(self,color)
self.face_attributes = utils.get_face_attributes_from_color(self,color)
def set_image_texture(self,texture_map,inplace=True):
mesh = self if inplace else copy.deepcopy(self)
if isinstance(texture_map,str):
texture_map = PIL.Image.open(texture_map)
texture_map = np.array(texture_map,dtype=np.float) / 255.0
texture_map = torch.tensor(texture_map,dtype=torch.float).to(device).permute(2,0,1).unsqueeze(0)
mesh.texture_map = texture_map
return mesh
def divide(self,inplace=True):
mesh = self if inplace else copy.deepcopy(self)
new_vertices, new_faces, new_vertex_normals, new_face_uvs, new_face_normals = utils.add_vertices(mesh)
mesh.vertices = new_vertices
mesh.faces = new_faces
mesh.vertex_normals = new_vertex_normals
mesh.face_uvs = new_face_uvs
mesh.face_normals = new_face_normals
return mesh
def divide_with_labels(self, face_label_map, inplace=True):
mesh = self if inplace else copy.deepcopy(self)
new_vertices, new_faces, new_vertex_normals, new_face_label_map = utils.add_vertices_with_labels(mesh, face_label_map)
mesh.vertices = new_vertices
mesh.faces = new_faces
mesh.vertex_normals = new_vertex_normals
return mesh, new_face_label_map
def export(self, file, extension="obj", color=None):
with open(file, "w+") as f:
if extension == "obj":
for vi, v in enumerate(self.vertices):
if color is None:
f.write("v %f %f %f\n" % (v[0], v[1], v[2]))
else:
f.write("v %f %f %f %f %f %f\n" % (v[0], v[1], v[2], color[vi][0], color[vi][1], color[vi][2]))
if self.vertex_normals is not None:
f.write("vn %f %f %f\n" % (self.vertex_normals[vi, 0], self.vertex_normals[vi, 1], self.vertex_normals[vi, 2]))
for face in self.faces:
f.write("f %d %d %d\n" % (face[0] + 1, face[1] + 1, face[2] + 1))
elif extension == "ply":
f.write("ply\nformat ascii 1.0\n")
f.write("element vertex {}\n".format(self.vertices.shape[0]))
f.write("property float32 x\nproperty float32 y\nproperty float32 z\n")
f.write("property uchar red\nproperty uchar green\nproperty uchar blue\n")
f.write("element face {}\n".format(self.faces.shape[0]))
f.write("property list uint8 int32 vertex_index\n")
f.write("end_header\n")
for vi, v in enumerate(self.vertices):
f.write("%f %f %f %d %d %d\n" % (v[0], v[1], v[2], color[vi][0], color[vi][1], color[vi][2]))
for face in self.faces:
f.write("3 %d %d %d\n" % (face[0], face[1], face[2]))
else:
print("error -- mesh.py: incorrect extension provdied in mesh.export()")