forked from GitBoSun/ARAPReg
-
Notifications
You must be signed in to change notification settings - Fork 2
/
Copy pathmeshdata.py
110 lines (97 loc) · 4.46 KB
/
meshdata.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
import openmesh as om
from datasets import SMAL, DFaust, Bone
from sklearn.decomposition import PCA
import numpy as np
import sys
sys.path.append("./")
import torch
import os
class MeshData(object):
def __init__(self,
root,
template_fp,
dataset='DFaust',
transform=None,
pre_transform=None,
pca_n_comp=8,
vert_pca=False,
heat_kernel=False):
self.root = root
self.template_fp = template_fp
self.dataset = dataset
self.transform = transform
self.pre_transform = pre_transform
self.train_dataset = None
self.test_dataste = None
self.template_points = None
self.template_face = None
self.mean = None
self.std = None
self.num_nodes = None
self.use_vert_pca = vert_pca
self.use_heat_kernel = heat_kernel
self.pca = PCA(n_components=pca_n_comp)
self.load()
def load(self):
if self.dataset=='SMAL':
self.train_dataset = SMAL(self.root,
train=True,
transform=self.transform,
pre_transform=self.pre_transform)
self.test_dataset = SMAL(self.root,
train=False,
transform=self.transform,
pre_transform=self.pre_transform)
elif self.dataset=='DFaust':
self.train_dataset = DFaust(self.root,
train=True,
transform=self.transform,
pre_transform=self.pre_transform)
self.test_dataset = DFaust(self.root,
train=False,
transform=self.transform,
pre_transform=self.pre_transform)
elif self.dataset=='Bone':
self.train_dataset = Bone(self.root,
train=True,
transform=self.transform,
pre_transform=self.pre_transform)
self.test_dataset = Bone(self.root,
train=False,
transform=self.transform,
pre_transform=self.pre_transform)
tmp_mesh = om.read_trimesh(self.template_fp)
self.template_points = tmp_mesh.points()
self.template_face = tmp_mesh.face_vertex_indices()
self.num_nodes = self.train_dataset[0].num_nodes
self.num_train_graph = len(self.train_dataset)
self.num_test_graph = len(self.test_dataset)
self.mean = self.train_dataset.data.x.view(self.num_train_graph, -1,
3).mean(dim=0)
self.std = self.train_dataset.data.x.view(self.num_train_graph, -1,
3).std(dim=0)
self.normalize()
def normalize(self):
vertices_train = self.train_dataset.data.x.view(self.num_train_graph, -1, 3).numpy()
vertices_test = self.test_dataset.data.x.view(self.num_test_graph, -1, 3).numpy()
print('Normalizing...')
self.train_dataset.data.x = (
(self.train_dataset.data.x.view(self.num_train_graph, -1, 3) -
self.mean) / self.std).view(-1, 3)
self.test_dataset.data.x = (
(self.test_dataset.data.x.view(self.num_test_graph, -1, 3) -
self.mean) / self.std).view(-1, 3)
if self.use_vert_pca:
print("Computing vertex PCA...")
self.pca.fit(np.reshape(vertices_train, (self.num_train_graph, -1)))
pca_axes = self.pca.components_
train_pca_sv= np.matmul(np.reshape(vertices_train, (self.num_train_graph, -1)), pca_axes.transpose())
test_pca_sv = np.matmul(np.reshape(vertices_test, (self.num_test_graph, -1)), pca_axes.transpose())
pca_sv_mean = np.mean(train_pca_sv, axis=0)
pca_sv_std = np.std(train_pca_sv, axis=0)
self.train_pca_sv = (train_pca_sv - pca_sv_mean)/pca_sv_std
self.test_pca_sv = (test_pca_sv - pca_sv_mean)/pca_sv_std
print('Done!')
def save_mesh(self, fp, x):
x = x * self.std + self.mean
om.write_mesh(fp, om.TriMesh(x.numpy(), self.template_face))