init commitHEADmaster

1 files changed, 116 insertions, 0 deletions

diff --git a/gaussian.py b/gaussian.py
new file mode 100644
index 0000000..10f23f3
--- /dev/null
+++ b/gaussian.py
@@ -0,0 +1,116 @@
+from utils import quat2rot
+import numpy as np
+
+class Gaussian:
+    def __init__(self, positions, scales, rotations, features_dc, features_rest, opacity):
+        self.num_gaussians = positions.shape[0]
+        self.positions = positions
+
+        self.scales = scales
+        self.rotations = rotations
+
+        # self.scales = np.exp(scales)
+        
+        # rot_length = np.repeat(np.linalg.norm(rotations, axis=1), 4).reshape(-1, 4)
+        # self.rotations = rotations / rot_length
+
+        self.features_dc = features_dc
+        self.features_rest = features_rest
+
+        self.opacity = opacity
+
+    def empty_with_cap(cap: int) -> "Gaussian":
+        positions = np.zeros((cap, 3))
+        scales = np.zeros((cap, 3))
+        rotations = np.zeros((cap, 4))
+        features_dc = np.zeros((cap, 1, 3))
+        features_rest = np.zeros((cap, 15, 3))
+        opacity = np.zeros((cap, 1))
+
+        gaussian = Gaussian(positions, scales, rotations, features_dc, features_rest, opacity)
+        gaussian.num_gaussians = 0
+        return gaussian
+
+    @property
+    def sh(self):
+        return np.concatenate((self.features_dc, self.features_rest), axis=1)
+    
+    def calc(self, x, id: int, scale: float):
+        x = x * scale
+        S = np.diag(self.scales[id])
+        R = quat2rot(self.rotations[id])
+        M = S @ R
+
+        Sigma = M.T @ M
+        Sigma = np.linalg.inv(Sigma)
+
+        return np.dot(x, Sigma @ x)
+        # M = R @ S
+        # return np.dot(x, M @ M.T @ x)
+
+    def clip_to_box(self, min: np.ndarray, max: np.ndarray):
+        indices = ((self.positions > min) & (self.positions < max)).all(axis=1)
+        self.apply_filter(indices)
+    
+    def apply_filter(self, filter):
+        self.positions = self.positions[filter]
+        self.scales = self.scales[filter]
+        self.rotations = self.rotations[filter]
+        self.features_dc = self.features_dc[filter]
+        self.features_rest = self.features_rest[filter]
+        self.opacity = self.opacity[filter]
+        self.num_gaussians = self.positions.shape[0]
+
+    def add(self, position, scale, rotation, features_dc, features_rest, opacity):
+        n = self.num_gaussians
+
+        self.positions[n] = position
+        self.scales[n] = scale
+        self.rotations[n] = rotation
+        self.features_dc[n] = features_dc
+        self.features_rest[n] = features_rest
+        self.opacity[n] = opacity
+
+        self.num_gaussians += 1
+        return n
+    
+    def concat(self, other: "Gaussian"):
+        self.positions = np.concatenate((self.positions, other.positions), axis=0)
+        self.scales = np.concatenate((self.scales, other.scales), axis=0)
+        self.rotations = np.concatenate((self.rotations, other.rotations), axis=0)
+        self.features_dc = np.concatenate((self.features_dc, other.features_dc), axis=0)
+        self.features_rest = np.concatenate((self.features_rest, other.features_rest), axis=0)
+        self.opacity = np.concatenate((self.opacity, other.opacity), axis=0)
+        self.num_gaussians = self.num_gaussians + other.num_gaussians
+
+    def replace(self, id1: int, id2: int,
+                position, scale, rotation, sh, opacity):
+        # 用新的 gaussian 替换 id1 位置的 gaussian
+        self.positions[id1] = position
+        self.scales[id1] = scale
+        self.rotations[id1] = rotation
+        self.features_dc[id1] = sh[0, :]
+        self.features_rest[id1] = sh[1:, :]
+        self.opacity[id1] = opacity
+
+        # 把最后一个 gaussian 放到 id2 的位置上
+        n = self.num_gaussians
+        self.positions[id2] = self.positions[n - 1]
+        self.scales[id2] = self.scales[n - 1]
+        self.rotations[id2] = self.rotations[n - 1]
+        self.features_dc[id2] = self.features_dc[n - 1]
+        self.features_rest[id2] = self.features_rest[n - 1]
+        self.opacity[id2] = self.opacity[n - 1]
+        
+        self.num_gaussians -= 1
+        pass
+
+    def copy(self) -> "Gaussian":
+        positions = np.copy(self.positions)
+        scales = np.copy(self.scales)
+        rotations = np.copy(self.rotations)
+        features_dc = np.copy(self.features_dc)
+        features_rest = np.copy(self.features_rest)
+        opacity = np.copy(self.opacity)
+
+        return Gaussian(positions, scales, rotations, features_dc, features_rest, opacity)
\ No newline at end of file