generated from dongliangcao/pytorch-framework
-
Notifications
You must be signed in to change notification settings - Fork 0
/
preprocess.py
70 lines (56 loc) · 2.54 KB
/
preprocess.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
import os
import scipy.io as sio
import numpy as np
from argparse import ArgumentParser
from glob import glob
from tqdm import tqdm
import torch
from utils.geometry_util import laplacian_decomposition, get_operators
from utils.shape_util import read_shape, compute_geodesic_distmat, write_off
if __name__ == '__main__':
# parse arguments
parser = ArgumentParser('Preprocess .off files')
parser.add_argument('--data_root', required=True, help='data root contains /off sub-folder.')
parser.add_argument('--n_eig', type=int, default=300, help='number of eigenvectors/values to compute.')
parser.add_argument('--no_eig', action='store_true', help='no laplacian eigen-decomposition')
parser.add_argument('--no_dist', action='store_true', help='no geodesic matrix.')
parser.add_argument('--no_normalize', action='store_true', help='no normalization of face area.')
args = parser.parse_args()
# sanity check
data_root = args.data_root
n_eig = args.n_eig
no_eig = args.no_eig
no_dist = args.no_dist
no_normalize = args.no_normalize
assert n_eig > 0, f'Invalid n_eig: {n_eig}'
assert os.path.isdir(data_root), f'Invalid data root: {data_root}'
if not no_eig:
spectral_dir = os.path.join(data_root, 'diffusion')
os.makedirs(spectral_dir, exist_ok=True)
if not no_dist:
dist_dir = os.path.join(data_root, 'dist')
os.makedirs(dist_dir, exist_ok=True)
# read .off files
off_files = sorted(glob(os.path.join(data_root, 'off', '*.off')))
assert len(off_files) != 0
for off_file in tqdm(off_files):
verts, faces = read_shape(off_file)
filename = os.path.basename(off_file)
if not no_normalize:
# center shape
verts -= np.mean(verts, axis=0)
# normalize verts by sqrt face area
old_sqrt_area = laplacian_decomposition(verts=verts, faces=faces, k=1)[-1]
print(f'Old face sqrt area: {old_sqrt_area:.3f}')
verts /= old_sqrt_area
# save new verts and faces
write_off(off_file, verts, faces)
if not no_eig:
# recompute laplacian decomposition
get_operators(torch.from_numpy(verts).float(), torch.from_numpy(faces).long(),
k=n_eig, cache_dir=spectral_dir)
if not no_dist:
# compute distance matrix
dist_mat = compute_geodesic_distmat(verts, faces)
# save results
sio.savemat(os.path.join(dist_dir, filename.replace('.off', '.mat')), {'dist': dist_mat})