-
Notifications
You must be signed in to change notification settings - Fork 3
/
gauss_handler.py
208 lines (151 loc) · 7.04 KB
/
gauss_handler.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
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
205
206
207
import numpy as np
import torch
import time
from math import floor
"""
These functions convert a gaussian scale and rotation into a covariance matrix.
Originally provided: https://github.com/graphdeco-inria/gaussian-splatting
"""
def strip_lowerdiag(L):
uncertainty = torch.zeros((L.shape[0], 6), dtype=torch.float, device="cuda")
uncertainty[:, 0] = L[:, 0, 0]
uncertainty[:, 1] = L[:, 0, 1]
uncertainty[:, 2] = L[:, 0, 2]
uncertainty[:, 3] = L[:, 1, 1]
uncertainty[:, 4] = L[:, 1, 2]
uncertainty[:, 5] = L[:, 2, 2]
return uncertainty
def strip_symmetric(sym):
return strip_lowerdiag(sym)
def build_rotation(q):
#norm = torch.sqrt(r[:,0]*r[:,0] + r[:,1]*r[:,1] + r[:,2]*r[:,2] + r[:,3]*r[:,3])
#q = r / norm[:, None]
R = torch.zeros((q.size(0), 3, 3), device='cuda')
r = q[:, 0]
x = q[:, 1]
y = q[:, 2]
z = q[:, 3]
R[:, 0, 0] = 1 - 2 * (y*y + z*z)
R[:, 0, 1] = 2 * (x*y - r*z)
R[:, 0, 2] = 2 * (x*z + r*y)
R[:, 1, 0] = 2 * (x*y + r*z)
R[:, 1, 1] = 1 - 2 * (x*x + z*z)
R[:, 1, 2] = 2 * (y*z - r*x)
R[:, 2, 0] = 2 * (x*z - r*y)
R[:, 2, 1] = 2 * (y*z + r*x)
R[:, 2, 2] = 1 - 2 * (x*x + y*y)
return R
def build_scaling_rotation(s, r):
L = torch.zeros((s.shape[0], 3, 3), dtype=torch.float, device="cuda")
R = build_rotation(r)
L[:,0,0] = torch.exp(s[:,0])
L[:,1,1] = torch.exp(s[:,1])
L[:,2,2] = torch.exp(s[:,2])
L = R @ L
return L
def build_covariance_from_scaling_rotation(scaling, scaling_modifier, rotation):
L = build_scaling_rotation(scaling_modifier * scaling, rotation)
actual_covariance = L @ L.transpose(1, 2)
return actual_covariance
class Gaussians():
"""
Manages all loaded gaussians in the renderer
"""
def __init__(self, xyz, scales, rots, colours, opacities):
self.xyz = xyz
self.scales = scales
self.rots = rots
self.opacities = opacities
self.colours = colours
self.scaling_modifier = 1.0
# Calculates 3D covariance matrices
covariances = build_covariance_from_scaling_rotation(scales, self.scaling_modifier, rots)
# Validate and add covariances matrices
self.validate_and_add_covariances(covariances)
def non_posdef_covariances(self, covariances):
"""
Returns a boolean mask of which covariances are definitepositive
"""
return torch.any(torch.linalg.eigvals(covariances).real <= 0, 1)
def clamp_covariances(self, covariances, mask=None, epsilon=1e-6):
"""
Clips Eigenvalues to positive to enforce covariances to be positive-definite
Credit: MultiTrickFox
"""
if mask is None:
mask = torch.ones(covariances.shape[0], dtype=torch.bool)
eigvals, eigvecs = torch.linalg.eigh(covariances[mask])
eigvals = torch.clamp(eigvals, min=epsilon)
covariances[mask] = eigvecs @ torch.diag_embed(eigvals) @ eigvecs.transpose(-1, -2)
return covariances
def regularise_covariances(self, covariances, mask=None, epsilon=1e-6):
"""
Increases the value of the diagonal of the covariances matrices to ensure it is positive-definite
"""
if mask is None:
mask = torch.ones(covariances.shape[0], dtype=torch.bool)
eye_matrix = epsilon * torch.eye(3, device=self.xyz.get_device()).expand(mask.sum(), 3, 3)
covariances[mask] += eye_matrix
return covariances
def validate_and_add_covariances(self, covariances):
"""
Regularises Gaussian covariances and ensures that all covariances are positive-definite
since sometimes gaussian values are slightly wrong when loaded (most likely due to floating point errors)
"""
# Regularises gaussians with a low factor, which ensures almost all covaricnes are positive-definite
covariances = self.regularise_covariances(covariances)
# Check if any non positive-definite covariances exist, if so, clamp gaussians to ensure
# all covariances are positive-definite
non_positive_covariances = self.non_posdef_covariances(covariances)
if non_positive_covariances.sum() > 0:
covariances = self.clamp_covariances(covariances, mask=non_positive_covariances, epsilon=0.001)
self.covariances = covariances
# If still not positive-definite then delete erroneous Gaussians
non_positive_covariances = self.non_posdef_covariances(covariances)
if non_positive_covariances.sum() > 0:
self.filter_gaussians(~non_positive_covariances)
def filter_gaussians(self, filter_indices):
"""
Filters gaussians based on given indices
"""
self.xyz = self.xyz[filter_indices]
self.scales = self.scales[filter_indices]
self.rots = self.rots[filter_indices]
self.colours = self.colours[filter_indices]
self.opacities = self.opacities[filter_indices]
self.covariances = self.covariances[filter_indices]
def apply_min_opacity(self, min_opacity):
"""
Removes gaussians with opacity lower than the min_opacity
"""
if min_opacity > 0.0:
valid_gaussians_indices = self.opacities > (min_opacity)
self.filter_gaussians(valid_gaussians_indices)
def apply_bounding_box(self, bounding_box_min, bounding_box_max):
"""
Removes gaussians outside of the bounding box
"""
valid_gaussians_indices = torch.logical_not(torch.zeros(self.xyz.shape[0], dtype=torch.bool, device=self.xyz.get_device()))
if bounding_box_min is not None:
bounding_box_min_indices = (self.xyz[:,0] > bounding_box_min[0]) & (self.xyz[:,1] > bounding_box_min[1]) \
& (self.xyz[:,2] > bounding_box_min[2])
valid_gaussians_indices = valid_gaussians_indices & bounding_box_min_indices
if bounding_box_max is not None:
bounding_box_max_indices = (self.xyz[:,0] < bounding_box_max[0]) & (self.xyz[:,1] < bounding_box_max[1]) \
& (self.xyz[:,2] < bounding_box_max[2])
valid_gaussians_indices = valid_gaussians_indices & bounding_box_max_indices
self.filter_gaussians(valid_gaussians_indices)
def cull_large_gaussians(self, cull_gauss_size_percent):
"""
Orders the gaussians by size and removes gaussians with a size greater than the 'cull_gauss_size_percent'
"""
gaussian_sizes = self.get_gaussian_sizes()
cull_index = floor(gaussian_sizes.shape[0] *(1-cull_gauss_size_percent))
sorted_sizes, sorted_indices = torch.sort(gaussian_sizes)
culled_gaussians = sorted_indices[:cull_index]
self.filter_gaussians(culled_gaussians)
def get_gaussian_sizes(self):
"""
Orders the gaussians by volume size after calculating the exponent (to prioritse larger Gaussians)
"""
return torch.sqrt(torch.sum(torch.pow(torch.exp(self.scales * self.scaling_modifier), 2), axis=1))