forked from PrimitiveAnything/PrimitiveAnything
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathmeshUtils.py
More file actions
204 lines (178 loc) · 13.3 KB
/
Copy pathmeshUtils.py
File metadata and controls
204 lines (178 loc) · 13.3 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
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
from modules.transformer import rotate_module, translate_module
import matplotlib.pyplot as plt
import torch
colormap = torch.Tensor([[0.000000, 0.000000, 0.515625], [0.000000, 0.000000, 0.531250], [0.000000, 0.000000, 0.546875],
[0.000000, 0.000000, 0.562500], [0.000000, 0.000000, 0.578125], [0.000000, 0.000000, 0.593750],
[0.000000, 0.000000, 0.609375], [0.000000, 0.000000, 0.625000], [0.000000, 0.000000, 0.640625],
[0.000000, 0.000000, 0.656250], [0.000000, 0.000000, 0.671875], [0.000000, 0.000000, 0.687500],
[0.000000, 0.000000, 0.703125], [0.000000, 0.000000, 0.718750], [0.000000, 0.000000, 0.734375],
[0.000000, 0.000000, 0.750000], [0.000000, 0.000000, 0.765625], [0.000000, 0.000000, 0.781250],
[0.000000, 0.000000, 0.796875], [0.000000, 0.000000, 0.812500], [0.000000, 0.000000, 0.828125],
[0.000000, 0.000000, 0.843750], [0.000000, 0.000000, 0.859375], [0.000000, 0.000000, 0.875000],
[0.000000, 0.000000, 0.890625], [0.000000, 0.000000, 0.906250], [0.000000, 0.000000, 0.921875],
[0.000000, 0.000000, 0.937500], [0.000000, 0.000000, 0.953125], [0.000000, 0.000000, 0.968750],
[0.000000, 0.000000, 0.984375], [0.000000, 0.000000, 1.000000], [0.000000, 0.015625, 1.000000],
[0.000000, 0.031250, 1.000000], [0.000000, 0.046875, 1.000000], [0.000000, 0.062500, 1.000000],
[0.000000, 0.078125, 1.000000], [0.000000, 0.093750, 1.000000], [0.000000, 0.109375, 1.000000],
[0.000000, 0.125000, 1.000000], [0.000000, 0.140625, 1.000000], [0.000000, 0.156250, 1.000000],
[0.000000, 0.171875, 1.000000], [0.000000, 0.187500, 1.000000], [0.000000, 0.203125, 1.000000],
[0.000000, 0.218750, 1.000000], [0.000000, 0.234375, 1.000000], [0.000000, 0.250000, 1.000000],
[0.000000, 0.265625, 1.000000], [0.000000, 0.281250, 1.000000], [0.000000, 0.296875, 1.000000],
[0.000000, 0.312500, 1.000000], [0.000000, 0.328125, 1.000000], [0.000000, 0.343750, 1.000000],
[0.000000, 0.359375, 1.000000], [0.000000, 0.375000, 1.000000], [0.000000, 0.390625, 1.000000],
[0.000000, 0.406250, 1.000000], [0.000000, 0.421875, 1.000000], [0.000000, 0.437500, 1.000000],
[0.000000, 0.453125, 1.000000], [0.000000, 0.468750, 1.000000], [0.000000, 0.484375, 1.000000],
[0.000000, 0.500000, 1.000000], [0.000000, 0.515625, 1.000000], [0.000000, 0.531250, 1.000000],
[0.000000, 0.546875, 1.000000], [0.000000, 0.562500, 1.000000], [0.000000, 0.578125, 1.000000],
[0.000000, 0.593750, 1.000000], [0.000000, 0.609375, 1.000000], [0.000000, 0.625000, 1.000000],
[0.000000, 0.640625, 1.000000], [0.000000, 0.656250, 1.000000], [0.000000, 0.671875, 1.000000],
[0.000000, 0.687500, 1.000000], [0.000000, 0.703125, 1.000000], [0.000000, 0.718750, 1.000000],
[0.000000, 0.734375, 1.000000], [0.000000, 0.750000, 1.000000], [0.000000, 0.765625, 1.000000],
[0.000000, 0.781250, 1.000000], [0.000000, 0.796875, 1.000000], [0.000000, 0.812500, 1.000000],
[0.000000, 0.828125, 1.000000], [0.000000, 0.843750, 1.000000], [0.000000, 0.859375, 1.000000],
[0.000000, 0.875000, 1.000000], [0.000000, 0.890625, 1.000000], [0.000000, 0.906250, 1.000000],
[0.000000, 0.921875, 1.000000], [0.000000, 0.937500, 1.000000], [0.000000, 0.953125, 1.000000],
[0.000000, 0.968750, 1.000000], [0.000000, 0.984375, 1.000000], [0.000000, 1.000000, 1.000000],
[0.015625, 1.000000, 0.984375], [0.031250, 1.000000, 0.968750], [0.046875, 1.000000, 0.953125],
[0.062500, 1.000000, 0.937500], [0.078125, 1.000000, 0.921875], [0.093750, 1.000000, 0.906250],
[0.109375, 1.000000, 0.890625], [0.125000, 1.000000, 0.875000], [0.140625, 1.000000, 0.859375],
[0.156250, 1.000000, 0.843750], [0.171875, 1.000000, 0.828125], [0.187500, 1.000000, 0.812500],
[0.203125, 1.000000, 0.796875], [0.218750, 1.000000, 0.781250], [0.234375, 1.000000, 0.765625],
[0.250000, 1.000000, 0.750000], [0.265625, 1.000000, 0.734375], [0.281250, 1.000000, 0.718750],
[0.296875, 1.000000, 0.703125], [0.312500, 1.000000, 0.687500], [0.328125, 1.000000, 0.671875],
[0.343750, 1.000000, 0.656250], [0.359375, 1.000000, 0.640625], [0.375000, 1.000000, 0.625000],
[0.390625, 1.000000, 0.609375], [0.406250, 1.000000, 0.593750], [0.421875, 1.000000, 0.578125],
[0.437500, 1.000000, 0.562500], [0.453125, 1.000000, 0.546875], [0.468750, 1.000000, 0.531250],
[0.484375, 1.000000, 0.515625], [0.500000, 1.000000, 0.500000], [0.515625, 1.000000, 0.484375],
[0.531250, 1.000000, 0.468750], [0.546875, 1.000000, 0.453125], [0.562500, 1.000000, 0.437500],
[0.578125, 1.000000, 0.421875], [0.593750, 1.000000, 0.406250], [0.609375, 1.000000, 0.390625],
[0.625000, 1.000000, 0.375000], [0.640625, 1.000000, 0.359375], [0.656250, 1.000000, 0.343750],
[0.671875, 1.000000, 0.328125], [0.687500, 1.000000, 0.312500], [0.703125, 1.000000, 0.296875],
[0.718750, 1.000000, 0.281250], [0.734375, 1.000000, 0.265625], [0.750000, 1.000000, 0.250000],
[0.765625, 1.000000, 0.234375], [0.781250, 1.000000, 0.218750], [0.796875, 1.000000, 0.203125],
[0.812500, 1.000000, 0.187500], [0.828125, 1.000000, 0.171875], [0.843750, 1.000000, 0.156250],
[0.859375, 1.000000, 0.140625], [0.875000, 1.000000, 0.125000], [0.890625, 1.000000, 0.109375],
[0.906250, 1.000000, 0.093750], [0.921875, 1.000000, 0.078125], [0.937500, 1.000000, 0.062500],
[0.953125, 1.000000, 0.046875], [0.968750, 1.000000, 0.031250], [0.984375, 1.000000, 0.015625],
[1.000000, 1.000000, 0.000000], [1.000000, 0.984375, 0.000000], [1.000000, 0.968750, 0.000000],
[1.000000, 0.953125, 0.000000], [1.000000, 0.937500, 0.000000], [1.000000, 0.921875, 0.000000],
[1.000000, 0.906250, 0.000000], [1.000000, 0.890625, 0.000000], [1.000000, 0.875000, 0.000000],
[1.000000, 0.859375, 0.000000], [1.000000, 0.843750, 0.000000], [1.000000, 0.828125, 0.000000],
[1.000000, 0.812500, 0.000000], [1.000000, 0.796875, 0.000000], [1.000000, 0.781250, 0.000000],
[1.000000, 0.765625, 0.000000], [1.000000, 0.750000, 0.000000], [1.000000, 0.734375, 0.000000],
[1.000000, 0.718750, 0.000000], [1.000000, 0.703125, 0.000000], [1.000000, 0.687500, 0.000000],
[1.000000, 0.671875, 0.000000], [1.000000, 0.656250, 0.000000], [1.000000, 0.640625, 0.000000],
[1.000000, 0.625000, 0.000000], [1.000000, 0.609375, 0.000000], [1.000000, 0.593750, 0.000000],
[1.000000, 0.578125, 0.000000], [1.000000, 0.562500, 0.000000], [1.000000, 0.546875, 0.000000],
[1.000000, 0.531250, 0.000000], [1.000000, 0.515625, 0.000000], [1.000000, 0.500000, 0.000000],
[1.000000, 0.484375, 0.000000], [1.000000, 0.468750, 0.000000], [1.000000, 0.453125, 0.000000],
[1.000000, 0.437500, 0.000000], [1.000000, 0.421875, 0.000000], [1.000000, 0.406250, 0.000000],
[1.000000, 0.390625, 0.000000], [1.000000, 0.375000, 0.000000], [1.000000, 0.359375, 0.000000],
[1.000000, 0.343750, 0.000000], [1.000000, 0.328125, 0.000000], [1.000000, 0.312500, 0.000000],
[1.000000, 0.296875, 0.000000], [1.000000, 0.281250, 0.000000], [1.000000, 0.265625, 0.000000],
[1.000000, 0.250000, 0.000000], [1.000000, 0.234375, 0.000000], [1.000000, 0.218750, 0.000000],
[1.000000, 0.203125, 0.000000], [1.000000, 0.187500, 0.000000], [1.000000, 0.171875, 0.000000],
[1.000000, 0.156250, 0.000000], [1.000000, 0.140625, 0.000000], [1.000000, 0.125000, 0.000000],
[1.000000, 0.109375, 0.000000], [1.000000, 0.093750, 0.000000], [1.000000, 0.078125, 0.000000],
[1.000000, 0.062500, 0.000000], [1.000000, 0.046875, 0.000000], [1.000000, 0.031250, 0.000000],
[1.000000, 0.015625, 0.000000], [1.000000, 0.000000, 0.000000], [0.984375, 0.000000, 0.000000],
[0.968750, 0.000000, 0.000000], [0.953125, 0.000000, 0.000000], [0.937500, 0.000000, 0.000000],
[0.921875, 0.000000, 0.000000], [0.906250, 0.000000, 0.000000], [0.890625, 0.000000, 0.000000],
[0.875000, 0.000000, 0.000000], [0.859375, 0.000000, 0.000000], [0.843750, 0.000000, 0.000000],
[0.828125, 0.000000, 0.000000], [0.812500, 0.000000, 0.000000], [0.796875, 0.000000, 0.000000],
[0.781250, 0.000000, 0.000000], [0.765625, 0.000000, 0.000000], [0.750000, 0.000000, 0.000000],
[0.734375, 0.000000, 0.000000], [0.718750, 0.000000, 0.000000], [0.703125, 0.000000, 0.000000],
[0.687500, 0.000000, 0.000000], [0.671875, 0.000000, 0.000000], [0.656250, 0.000000, 0.000000],
[0.640625, 0.000000, 0.000000], [0.625000, 0.000000, 0.000000], [0.609375, 0.000000, 0.000000],
[0.593750, 0.000000, 0.000000], [0.578125, 0.000000, 0.000000], [0.562500, 0.000000, 0.000000],
[0.546875, 0.000000, 0.000000], [0.531250, 0.000000, 0.000000], [0.515625, 0.000000, 0.000000],
[0.500000, 0.000000, 0.000000]])
cubeV = torch.Tensor(
[[0.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 1.0, 0.0], [0.0, 1.0, 1.0], [1.0, 0.0, 0.0], [1.0, 0.0, 1.0],
[1.0, 1.0, 0.0], [1.0, 1.0, 1.0]])
cubeV = 2 * cubeV - 1
cubeF = torch.Tensor(
[[1, 7, 5], [1, 3, 7], [1, 4, 3], [1, 2, 4], [3, 8, 7], [3, 4, 8], [5, 7, 8], [5, 8, 6], [1, 5, 6], [1, 6, 2],
[2, 6, 8], [2, 8, 4]])
import pdb
def cuboidMesh(shape):
# Shape is a Float Tensor
verts = cubeV.clone()
for d in range(0, 3):
verts[:, d] = verts[:, d] * shape[d]
return verts, cubeF.clone()
def shapeMesh(shape):
return cuboidMesh(shape)
def shapeVolume(shape):
return 8 * shape[0] * shape[1] * shape[2]
def unpack(part):
return (part[0:3], part[3:6], part[6:10])
def plot_voxels(voxels, block=True, color='r'):
voxels2 = voxels.__ge__(0.5)
cube_len = voxels.shape[0]
x, y, z = voxels2.nonzero()
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
azm = -97
ele = 112
ax.view_init(elev=ele, azim=azm)
ax.set_xlim([0, cube_len])
ax.set_ylim([0, cube_len])
ax.set_zlim([0, cube_len])
ax.scatter(x, y, z, zdir='z', c=color, s=20, marker='s')
plt.show(block=block)
def partVerticesAndFaces(predParts):
part_vertices = []
nParts = len(predParts)
for p in range(nParts):
shape, translation, quat = unpack(predParts[p])
verts, faces = shapeMesh(shape)
nVerts = verts.size(0)
verts = verts.view(1, nVerts, 3)
quat = quat.clone()
quat[0] = -quat[0]
quat = quat.view(1, 1, -1)
trans = translation.view(1, 1, -1)
verts = rotate_module(verts, quat)
verts = translate_module(verts, trans)
verts = verts.squeeze(0)
part_vertices.append(verts)
return part_vertices, faces
import pdb
def saveParts(predParts, outputFile, partIndsSpecific=[]):
# outputFile is obj
mtlfile = outputFile.replace("obj", "mtl")
foutMtl = open(mtlfile, 'w')
fout = open(outputFile, 'w')
mtlfile = mtlfile.split("/")
mtlfile = mtlfile[-1]
partCmaps = []
nParts = len(predParts)
for p in range(nParts):
partCmaps.append(colormap[p * (256 // nParts)])
foutMtl.write('newmtl m{}\nKd {} {} {}\nKa 0 0 0\n'.format(p, partCmaps[p][0], partCmaps[p][1], partCmaps[p][2]))
foutMtl.close()
fout.write('mtllib {}\n'.format(mtlfile))
if len(partIndsSpecific) == 0:
for i in range(len(predParts)):
partIndsSpecific.append(i)
vertsOffset = 0
for i in partIndsSpecific:
p = partIndsSpecific[i]
if predParts[p][10] > 0.5:
verts, faces = partVerticesAndFaces([predParts[p]])
verts = verts[0]
faces = faces + vertsOffset
fout.write('usemtl m{}\n'.format(p))
for vx in range(verts.size(0)):
fout.write('v {} {} {}\n'.format(verts[vx][0], verts[vx][1], verts[vx][2]))
for fx in range(faces.size(0)):
fout.write('f {} {} {}\n'.format(faces[fx][0], faces[fx][1], faces[fx][2]))
vertsOffset = vertsOffset + verts.size(0)
fout.close()
def savePredParts(predpart, outputfile):
pred_b = []
nParts = predpart.size(0)
for px in range(nParts):
pred_b.append(predpart[px, :].clone().data.cpu())
saveParts(pred_b, '{}'.format(outputfile))