forked from tneumann/cmm
-
Notifications
You must be signed in to change notification settings - Fork 0
/
compute_cmm.py
118 lines (100 loc) · 4.32 KB
/
compute_cmm.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
114
115
116
117
118
import plac
from os import path
import numpy as np
from scipy import sparse
from scipy.io import savemat
from cmmlib.inout import load_mesh, save_coff
from cmmlib import cmm
@plac.annotations(
K=('number of CMHBs', 'positional', None, int),
mu=('sparsity parameter mu', 'positional', None, float),
visualize=('visualize the weights?', 'flag', 'v'),
scaled=('respect triangle scaling?', 'flag', 's'),
output_dir=('output directory', 'option', 'o'),
maxiter=('maximum number of iterations', 'option', None, int),
ply=('output ply file?', 'flag', None),
off=('output off file?', 'flag', None)
)
def main(input_filename, K, mu, output_dir=None, visualize=False, scaled=False,
maxiter=None, ply=False, off=False):
if (off or ply) and not output_dir:
print "please specify an output directory"
return 1
if output_dir and not path.exists(output_dir):
print "%s does not exist" % output_dir
return 2
verts, tris = load_mesh(input_filename, normalize=True)
print "%d vertices, %d faces" % (len(verts), len(tris))
Phi_cpr, D = cmm.compressed_manifold_modes(
verts, tris, K, mu=mu, scaled=scaled,
maxiter=maxiter, verbose=100, return_D=True)
if D is None:
D_diag = np.ones(len(verts))
D = sparse.eye(len(verts))
else:
D_diag = D.data
if visualize:
from cmmlib.vis.weights import show_weights
show_weights(verts, tris, Phi_cpr)
if output_dir:
# save in simple text format
np.savetxt(path.join(output_dir, 'phi.txt'), Phi_cpr, fmt='%f')
np.savetxt(path.join(output_dir, 'D_diag.txt'), D_diag, fmt='%f')
# save in matlab format
savemat(path.join(output_dir, 'phi.mat'),
dict(verts=verts, tris=tris+1, phi=Phi_cpr, D=D))
# save HDF5 format if possible
try:
import h5py
except ImportError:
print "Cannot save as HDF5, please install the h5py module"
else:
with h5py.File(path.join(output_dir, 'phi.h5'), 'w') as f:
f['verts'] = verts
f['tris'] = tris
f['phi'] = Phi_cpr
f['d_diag'] = D_diag
# save NPY format
np.save(path.join(output_dir, 'phi.npy'), Phi_cpr)
np.save(path.join(output_dir, 'D_diag.npy'), Phi_cpr)
if off or ply:
# map phi scalars to colors
from mayavi.core.lut_manager import LUTManager
from cmmlib.vis.weights import _centered
lut = LUTManager(lut_mode='RdBu').lut.table.to_array()[:, :3]
colors = [
lut[(_centered(Phi_cpr[:, k]) * (lut.shape[0]-1)).astype(int)]
for k in xrange(K)]
# save in a single scene as a collage
w = int(np.ceil(np.sqrt(K))) if K > 6 else K
spacing = 1.2 * verts.ptp(axis=0)
all_verts = [verts + spacing * (1.5, 0, 0)]
all_tris = [tris]
all_color = [np.zeros(verts.shape, np.int) + 127]
for k in xrange(K):
all_verts.append(verts + spacing * (-(k % w), 0, int(k / w)))
all_tris.append(tris + len(verts) * (k+1))
all_color.append(colors[k])
if off:
save_coff(path.join(output_dir, 'input.off'),
verts.astype(np.float32), tris)
for k in xrange(K):
save_coff(path.join(output_dir, 'cmh_%03d.off' % k),
verts.astype(np.float32), tris, colors[k])
save_coff(path.join(output_dir, 'all.off'),
np.vstack(all_verts), np.vstack(all_tris),
np.vstack(all_color))
if ply:
from tvtk.api import tvtk
pd = tvtk.PolyData(
points=np.vstack(all_verts).astype(np.float32),
polys=np.vstack(all_tris).astype(np.uint32))
pd.point_data.scalars = np.vstack(all_color).astype(np.uint8)
pd.point_data.scalars.name = 'colors'
ply = tvtk.PLYWriter(
file_name=path.join(output_dir, 'all.ply'),
input=pd, color=(1, 1, 1))
ply.array_name = 'colors'
ply.write()
if __name__ == '__main__':
plac.call(main)