forked from limacv/GaussianSplattingViewer
-
Notifications
You must be signed in to change notification settings - Fork 3
/
Copy pathutil_gau.py
114 lines (101 loc) · 4.02 KB
/
util_gau.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
import numpy as np
from plyfile import PlyData
from dataclasses import dataclass
@dataclass
class GaussianData:
xyz: np.ndarray
rot: np.ndarray
scale: np.ndarray
opacity: np.ndarray
sh: np.ndarray
def flat(self) -> np.ndarray:
ret = np.concatenate([self.xyz, self.rot, self.scale, self.opacity, self.sh], axis=-1)
return np.ascontiguousarray(ret)
def __len__(self):
return len(self.xyz)
@property
def sh_dim(self):
return self.sh.shape[-1]
def naive_gaussian():
gau_xyz = np.array([
0, 0, 0,
1, 0, 0,
0, 1, 0,
0, 0, 1,
]).astype(np.float32).reshape(-1, 3)
gau_rot = np.array([
1, 0, 0, 0,
1, 0, 0, 0,
1, 0, 0, 0,
1, 0, 0, 0
]).astype(np.float32).reshape(-1, 4)
gau_s = np.array([
0.03, 0.03, 0.03,
0.2, 0.03, 0.03,
0.03, 0.2, 0.03,
0.03, 0.03, 0.2
]).astype(np.float32).reshape(-1, 3)
gau_c = np.array([
1, 0, 1,
1, 0, 0,
0, 1, 0,
0, 0, 1,
]).astype(np.float32).reshape(-1, 3)
gau_c = (gau_c - 0.5) / 0.28209
gau_a = np.array([
1, 1, 1, 1
]).astype(np.float32).reshape(-1, 1)
return GaussianData(
gau_xyz,
gau_rot,
gau_s,
gau_a,
gau_c
)
def load_ply(path):
max_sh_degree = 1
plydata = PlyData.read(path)
xyz = np.stack((np.asarray(plydata.elements[0]["x"]),
np.asarray(plydata.elements[0]["y"]),
np.asarray(plydata.elements[0]["z"])), axis=1)
opacities = np.asarray(plydata.elements[0]["opacity"])[..., np.newaxis]
features_dc = np.zeros((xyz.shape[0], 3, 1))
features_dc[:, 0, 0] = np.asarray(plydata.elements[0]["f_dc_0"])
features_dc[:, 1, 0] = np.asarray(plydata.elements[0]["f_dc_1"])
features_dc[:, 2, 0] = np.asarray(plydata.elements[0]["f_dc_2"])
extra_f_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("f_rest_")]
extra_f_names = sorted(extra_f_names, key = lambda x: int(x.split('_')[-1]))
# assert len(extra_f_names)==3 * (max_sh_degree + 1) ** 2 - 3
max_sh_degree = int(np.sqrt((len(extra_f_names)+3)//3)) - 1
features_extra = np.zeros((xyz.shape[0], len(extra_f_names)))
for idx, attr_name in enumerate(extra_f_names):
features_extra[:, idx] = np.asarray(plydata.elements[0][attr_name])
# Reshape (P,F*SH_coeffs) to (P, F, SH_coeffs except DC)
features_extra = features_extra.reshape((features_extra.shape[0], 3, (max_sh_degree + 1) ** 2 - 1))
features_extra = np.transpose(features_extra, [0, 2, 1])
scale_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("scale_")]
scale_names = sorted(scale_names, key = lambda x: int(x.split('_')[-1]))
scales = np.zeros((xyz.shape[0], len(scale_names)))
for idx, attr_name in enumerate(scale_names):
scales[:, idx] = np.asarray(plydata.elements[0][attr_name])
rot_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("rot")]
rot_names = sorted(rot_names, key = lambda x: int(x.split('_')[-1]))
rots = np.zeros((xyz.shape[0], len(rot_names)))
for idx, attr_name in enumerate(rot_names):
rots[:, idx] = np.asarray(plydata.elements[0][attr_name])
# pass activate function
xyz = xyz.astype(np.float32)
rots = rots / np.linalg.norm(rots, axis=-1, keepdims=True)
rots = rots.astype(np.float32)
scales = np.exp(scales)
scales = scales.astype(np.float32)
opacities = 1/(1 + np.exp(- opacities)) # sigmoid
opacities = opacities.astype(np.float32)
shs = np.concatenate([features_dc.reshape(-1, 3),
features_extra.reshape(len(features_dc), -1)], axis=-1).astype(np.float32)
shs = shs.astype(np.float32)
return GaussianData(xyz, rots, scales, opacities, shs)
if __name__ == "__main__":
gs = load_ply("C:\\Users\\MSI_NB\\Downloads\\viewers\\models\\train\\point_cloud\\iteration_7000\\point_cloud.ply")
a = gs.flat()
print(a.shape)