init commitHEADmaster

9 files changed, 929 insertions, 0 deletions

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..0e4f2a3
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,3 @@
+__pycache__
+.vscode
+*.pkl
diff --git a/aabb.py b/aabb.py
new file mode 100644
index 0000000..d913143
--- /dev/null
+++ b/aabb.py
@@ -0,0 +1,102 @@
+import numpy as np
+from utils import quat2rot
+
+def aabb_size(scale, rotation, scale_factor=1):
+    scale = scale_factor * scale
+    if scale.shape != (3, 3):
+        scale_matrix = np.diag(scale)
+    else:
+        scale_matrix = scale
+    
+    if rotation.shape != (3, 3):
+        rotation = rotation / np.linalg.norm(rotation)
+        rotation_matrix = quat2rot(rotation)
+        # r, x, y, z = rotation[0], rotation[1], rotation[2], rotation[3]
+        # rotation_matrix = np.array([
+        #     [1. - 2. * (y * y + z * z), 2. * (x * y - r * z), 2. * (x * z + r * y)],
+        #     [2. * (x * y + r * z), 1. - 2. * (x * x + z * z), 2. * (y * z - r * x)],
+        #     [2. * (x * z - r * y), 2. * (y * z + r * x), 1. - 2. * (x * x + y * y)]
+        # ], dtype=np.float32)
+    else:
+        rotation_matrix = rotation
+
+    M = scale_matrix @ rotation_matrix
+    R = M.T @ M
+
+    x_span = np.sqrt(R[0][0])
+    y_span = np.sqrt(R[1][1])
+    z_span = np.sqrt(R[2][2])
+
+    return x_span, y_span, z_span
+
+def process_aabb(position, scale, rotation, scale_factor=1):
+    x_span, y_span, z_span = aabb_size(scale, rotation, scale_factor)
+
+    xmin = position[0] - x_span
+    xmax = position[0] + x_span
+    ymin = position[1] - y_span
+    ymax = position[1] + y_span
+    zmin = position[2] - z_span
+    zmax = position[2] + z_span
+    
+    aabb = np.array([xmin, xmax, ymin, ymax, zmin, zmax])
+    if np.isnan(aabb).any():
+        print("nan detected!")
+        print(scale, rotation, aabb)
+
+    return aabb
+
+def solve_scale(rotation, aabb):
+    dest = np.array([
+        (aabb[1] - aabb[0]) * (aabb[1] - aabb[0]) / 4,
+        (aabb[3] - aabb[2]) * (aabb[3] - aabb[2]) / 4,
+        (aabb[5] - aabb[4]) * (aabb[5] - aabb[4]) / 4,
+    ])
+
+    rotation = rotation / np.linalg.norm(rotation)
+    r, x, y, z = rotation[0], rotation[1], rotation[2], rotation[3]
+    rotation_matrix = np.array([
+        [1. - 2. * (y * y + z * z), 2. * (x * y - r * z), 2. * (x * z + r * y)],
+        [2. * (x * y + r * z), 1. - 2. * (x * x + z * z), 2. * (y * z - r * x)],
+        [2. * (x * z - r * y), 2. * (y * z + r * x), 1. - 2. * (x * x + y * y)]
+    ], dtype=np.float32)
+
+    coeff = rotation_matrix.T * rotation_matrix.T
+    solution = np.linalg.solve(coeff, dest)
+    print(solution)
+    return np.sqrt(solution)
+
+
+def intersect(aabb1, aabb2):
+    new_aabb = np.zeros_like(aabb1)
+    new_aabb[0] = max(aabb1[0], aabb2[0])
+    new_aabb[1] = min(aabb1[1], aabb2[1])
+    new_aabb[2] = max(aabb1[2], aabb2[2])
+    new_aabb[3] = min(aabb1[3], aabb2[3])
+    new_aabb[4] = max(aabb1[4], aabb2[4])
+    new_aabb[5] = min(aabb1[5], aabb2[5])
+
+    if new_aabb[0] < new_aabb[1] and new_aabb[2] < new_aabb[3] and new_aabb[4] < new_aabb[5]:
+        return new_aabb
+    else:
+        return None
+
+def aabb_merge(aabb1, aabb2):
+    new_aabb = np.zeros_like(aabb1)
+    new_aabb[0] = min(aabb1[0], aabb2[0])
+    new_aabb[1] = max(aabb1[1], aabb2[1])
+    new_aabb[2] = min(aabb1[2], aabb2[2])
+    new_aabb[3] = max(aabb1[3], aabb2[3])
+    new_aabb[4] = min(aabb1[4], aabb2[4])
+    new_aabb[5] = max(aabb1[5], aabb2[5])
+    return new_aabb
+
+def center(aabb):
+    return np.array([
+        aabb[0] + aabb[1],
+        aabb[2] + aabb[3],
+        aabb[4] + aabb[5]
+    ]) / 2
+
+def min_interval(aabb):
+    return min(aabb[1] - aabb[0], min(aabb[3] - aabb[2], aabb[5] - aabb[4]))
\ No newline at end of file
diff --git a/dataloader.py b/dataloader.py
new file mode 100644
index 0000000..1ec880a
--- /dev/null
+++ b/dataloader.py
@@ -0,0 +1,116 @@
+from plyfile import PlyData, PlyElement
+from gaussian import Gaussian
+import numpy as np
+import os
+
+def load_ply(path: str, max_sh_degree: int) -> Gaussian:
+    ply_data = PlyData.read(path)
+
+    # position
+    xyz = np.stack(
+        (
+            np.asarray(ply_data.elements[0]["x"]),
+            np.asarray(ply_data.elements[0]["y"]),
+            np.asarray(ply_data.elements[0]["z"]),
+        ),
+        axis=1,
+    )
+
+    # opacity
+    def sigmoid(z):
+        return 1 / (1 + np.exp(-z))
+
+    opacities = np.asarray(ply_data.elements[0]["opacity"])[..., np.newaxis]
+    ## 过激活函数
+    opacities = sigmoid(opacities)
+
+    # sh
+    features_dc = np.zeros((xyz.shape[0], 3, 1), dtype=np.float32)
+    features_dc[:, 0, 0] = np.asarray(ply_data.elements[0]["f_dc_0"])
+    features_dc[:, 1, 0] = np.asarray(ply_data.elements[0]["f_dc_1"])
+    features_dc[:, 2, 0] = np.asarray(ply_data.elements[0]["f_dc_2"])
+
+    extra_f_names = [
+        p.name for p in ply_data.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
+    features_extra = np.zeros((xyz.shape[0], len(extra_f_names)), dtype=np.float32)
+    for idx, attr_name in enumerate(extra_f_names):
+        features_extra[:, idx] = np.asarray(ply_data.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_dc = np.transpose(features_dc, (0, 2, 1))
+    features_extra = np.transpose(features_extra, (0, 2, 1))
+    ## 拼接得到完整的 sh
+    ## sh = np.transpose(np.concatenate((features_dc, features_extra), axis=2), (0, 2, 1))
+
+    # scale
+    scale_names = [
+        p.name for p in ply_data.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)), dtype=np.float32)
+    for idx, attr_name in enumerate(scale_names):
+        scales[:, idx] = np.asarray(ply_data.elements[0][attr_name])
+    ## 过激活函数
+    scales = np.exp(scales)
+
+    # rotation
+    rot_names = [
+        p.name for p in ply_data.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)), dtype=np.float32)
+    for idx, attr_name in enumerate(rot_names):
+        rots[:, idx] = np.asarray(ply_data.elements[0][attr_name])
+    ## 过激活函数
+    rot_length = np.linalg.norm(rots, axis=1)
+    rots = rots / np.expand_dims(rot_length, axis=1)
+
+    return Gaussian(xyz, scales, rots, features_dc, features_extra, opacities)
+    ## return (num_gaussians, xyz, scales, rots, sh, opacities)
+
+def save_ply(path: str, gaussian: Gaussian):
+    print(f"saving {path}")
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+
+    def construct_list_of_attributes(gaussian):
+        l = ['x', 'y', 'z', 'nx', 'ny', 'nz']
+        # All channels except the 3 DC
+        for i in range(gaussian.features_dc.shape[1] * gaussian.features_dc.shape[2]):
+            l.append('f_dc_{}'.format(i))
+        for i in range(gaussian.features_rest.shape[1] * gaussian.features_rest.shape[2]):
+            l.append('f_rest_{}'.format(i))
+        l.append('opacity')
+        for i in range(gaussian.scales.shape[1]):
+            l.append('scale_{}'.format(i))
+        for i in range(gaussian.rotations.shape[1]):
+            l.append('rot_{}'.format(i))
+        return l
+    
+    def inverse_sigmoid(x):
+        return np.log(x/(1-x))
+    
+    xyz = gaussian.positions
+    normals = np.zeros_like(xyz)
+    f_dc = np.transpose(gaussian.features_dc, (0, 2, 1))
+    f_dc = np.reshape(f_dc, (f_dc.shape[0], f_dc.shape[1] * f_dc.shape[2]))
+    f_rest = np.transpose(gaussian.features_rest, (0, 2, 1))
+    f_rest = np.reshape(f_rest, (f_rest.shape[0], f_rest.shape[1] * f_rest.shape[2]))
+    opacities = inverse_sigmoid(gaussian.opacity)
+    scale = np.log(gaussian.scales)
+    rotation = gaussian.rotations
+
+    dtype_full = [(attribute, 'f4') for attribute in construct_list_of_attributes(gaussian)]
+
+    attributes = np.concatenate((xyz, normals, f_dc, f_rest, opacities, scale, rotation), axis=1)[:gaussian.num_gaussians]
+
+    elements = np.empty(attributes.shape[0], dtype=dtype_full)
+    elements[:] = list(map(tuple, attributes))
+    el = PlyElement.describe(elements, 'vertex')
+    PlyData([el]).write(path)
+
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
diff --git a/main.py b/main.py
new file mode 100644
index 0000000..521f341
--- /dev/null
+++ b/main.py
@@ -0,0 +1,99 @@
+from dataloader import load_ply, save_ply
+from gaussian import Gaussian
+from aabb import process_aabb
+from similarity import similarity
+from scipy.spatial import KDTree
+from merge import merge, merge_2, merge_3, merge_4
+from tqdm import tqdm
+import argparse
+import numpy as np
+import pickle as pkl
+
+NUMBER=1000
+
+parser = argparse.ArgumentParser("main.py")
+parser.add_argument('-m', '--model_path', type=str, default="/home/chuyan/Documents/code/gaussian/models/train/point_cloud/iteration_50000/point_cloud.ply")
+parser.add_argument('-o', '--output_path', type=str, default="/home/chuyan/Documents/code/gaussian/models/train/point_cloud/iteration_60000/point_cloud.ply")
+parser.add_argument('-d', '--sh_degrees', type=int, default=3)
+
+args = parser.parse_args()
+
+gaussian_model = load_ply(args.model_path, args.sh_degrees)
+gaussian_extended = Gaussian.empty_with_cap(NUMBER)
+
+# gaussian_model.clip_to_box(
+#     min=np.array([-0.2, -0.7, -0.3]),
+#     max=np.array([0, -0.3, 0])
+# )
+
+
+# kd_tree = KDTree(gaussian_model.positions)
+# similarities = []
+# for i in tqdm(range(gaussian_model.num_gaussians)):
+#     # print(f"currently running {i}")
+
+#     scale = gaussian_model.scales[i]
+#     radius = max(max(scale[0], scale[1]), scale[2]) * 1.5
+
+#     points = kd_tree.query_ball_point(gaussian_model.positions[i], r=radius, workers=18)
+#     for j in points:
+#         if j < i:
+#             geo_score, color_score = similarity(gaussian_model, j, i)
+#             if geo_score * color_score != 0:
+#                 similarities.append((j, i, geo_score, color_score, ))
+
+# sim = sorted(similarities, key=lambda pair: pair[2] * pair[3], reverse=True)
+# print(f"pair number is {len(sim)}")
+
+# with open('similarities.pkl', 'wb') as f:
+#     pkl.dump(sim, f)
+
+with open('similarities.pkl', 'rb') as f:
+    sim = pkl.load(f)
+
+def metric_with_opacity(pair):
+    i, j, gs, cs = pair[0], pair[1], pair[2], pair[3]   
+    o1, o2 = gaussian_model.opacity[i], gaussian_model.opacity[j]
+    return gs * (cs ** 0.3) \
+        * np.exp(np.dot(gaussian_model.rotations[i], gaussian_model.rotations[j])) \
+        * min(o1/o2, o2/o1)
+
+sim = sorted(sim, key=metric_with_opacity, reverse=True)
+print(f"There are {len(sim)} pairs")
+
+cnt = 0
+filter = np.ones(gaussian_model.num_gaussians, dtype=bool)
+removed = []
+used = set()
+for idx, (i, j, gs, cs) in enumerate(sim):
+    if i in used or j in used:
+        continue
+    position, scale, rotation, features_dc, features_rest, opacity = merge_3(gaussian_model, i, j)
+    if position is not None:
+        
+        cnt += 1
+        # if cnt != NUMBER:
+        #     continue
+        id = gaussian_extended.add(position, scale, rotation, features_dc, features_rest, opacity)
+        # print(f"opacity: ({gaussian_model.opacity[i]},{gaussian_model.opacity[j]}) -> {opacity}")
+        filter[i] = False
+        filter[j] = False
+        used.add(i)
+        used.add(j)
+
+        if cnt == NUMBER:
+            break
+    
+deleted_gaussian = gaussian_model.copy()
+deleted_gaussian.apply_filter(np.logical_not(filter))
+# deleted_gaussian.opacity = np.ones_like(deleted_gaussian.opacity)
+# gaussian_extended.opacity = np.ones_like(gaussian_extended.opacity)
+
+gaussian_model.apply_filter(filter)
+gaussian_model.concat(gaussian_extended)
+
+deleted_path = args.output_path.replace('60000', '60001')
+added_path = args.output_path.replace('60000', '60002')
+save_ply(args.output_path, gaussian_model)
+save_ply(deleted_path, deleted_gaussian)
+save_ply(added_path, gaussian_extended)
diff --git a/merge.py b/merge.py
new file mode 100644
index 0000000..f00ca1a
--- /dev/null
+++ b/merge.py
@@ -0,0 +1,259 @@
+from gaussian import Gaussian
+from aabb import process_aabb, aabb_merge, center, solve_scale, aabb_size
+from utils import quat2rot, rot2quat, sq2cov
+from scipy.optimize import minimize
+import numpy as np
+import random
+import time
+
+def merge_4(gaussian: Gaussian, id1: int, id2: int):
+    scale_1 = gaussian.scales[id1]
+    scale_2 = gaussian.scales[id2]
+    weight_1 = gaussian.opacity[id1] * np.prod(scale_1)
+    weight_2 = gaussian.opacity[id2] * np.prod(scale_2)
+    weight = weight_1 + weight_2
+    
+    position = (weight_1 * gaussian.positions[id1] + weight_2 * gaussian.positions[id2]) / weight
+    cov_1 = sq2cov(scale_1, gaussian.rotations[id1])
+    cov_2 = sq2cov(scale_2, gaussian.rotations[id2])
+    cov = (weight_1 * cov_1 + weight_2 * cov_2) / weight
+
+    eigenvalues, eigenvectors = np.linalg.eig(cov)
+    scale = np.sqrt(eigenvalues)
+    rotation = rot2quat(eigenvectors.T)
+
+    sh = (gaussian.sh[id1] + gaussian.sh[id2]) / 2.
+    features_dc = sh[0, :]
+    features_rest = sh[1:, :]
+    opacity = (weight_1 * gaussian.opacity[id1] + weight_2 * gaussian.opacity[id2]) / weight
+
+    return position, scale, rotation, features_dc, features_rest, opacity
+
+def merge_3(gaussian: Gaussian, id1: int, id2: int) -> float:
+    def make_gaussian(mu, s, r, o):
+        def calc_inner(pos: np.ndarray) -> float:
+            x = pos - mu
+            S = np.diag(s)
+            R = quat2rot(r)
+            M = S @ R
+            Sigma = np.linalg.inv(M.T @ M)
+            return o * np.exp(-0.5 * np.dot(x, Sigma @ x))
+        return calc_inner
+    opacity_f = gaussian.opacity[id1]
+    opacity_g = gaussian.opacity[id2]
+    mu_f = gaussian.positions[id1]
+    mu_g = gaussian.positions[id2]
+    cov_f = sq2cov(gaussian.scales[id1], gaussian.rotations[id1])
+    cov_g = sq2cov(gaussian.scales[id2], gaussian.rotations[id2])
+    sqrt_det_cov_f = np.sqrt(np.linalg.det(cov_f))
+    sqrt_det_cov_g = np.sqrt(np.linalg.det(cov_g))
+
+    def c_without_2pi(mu1, mu2, cov1, cov2):
+        mu = mu1 - mu2
+        cov = cov1 + cov2
+        return np.exp(-0.5 * np.dot(mu, cov @ mu)) / np.sqrt(np.linalg.det(cov))
+
+    f2integral = opacity_f * opacity_f * sqrt_det_cov_f
+    g2integral = opacity_g * opacity_g * sqrt_det_cov_g
+    fgintegral = (2 ** 1.5) * opacity_f * opacity_g * sqrt_det_cov_f * sqrt_det_cov_g * c_without_2pi(mu_f, mu_g, cov_f, cov_g)
+    
+    gaussian_f = make_gaussian(gaussian.positions[id1], gaussian.scales[id1], gaussian.rotations[id1], opacity_f)
+    gaussian_g = make_gaussian(gaussian.positions[id2], gaussian.scales[id2], gaussian.rotations[id2], opacity_g)
+
+    mu_0, s_0, r_0, o_0, point_min, basis = merge_geo(gaussian, id1, id2)
+    x_0 = np.concatenate((mu_0, s_0, r_0, o_0), axis=0)
+    
+    basis_x = basis[:, 0]
+    basis_y = basis[:, 1]
+    basis_z = basis[:, 2]
+
+    def target_inner_analytical(features) -> float:
+        mu = features[:3]
+        s = features[3:6]
+        r = features[6:10]
+        o = features[10]
+        cov = sq2cov(s, r)
+        sqrt_det_cov = np.sqrt(np.linalg.det(cov))
+
+        h2integral = o * o * sqrt_det_cov
+        fhintegral = (2 ** 1.5) * opacity_f * o * sqrt_det_cov_f * sqrt_det_cov * c_without_2pi(mu_f, mu, cov_f, cov)
+        ghintegral = (2 ** 1.5) * opacity_g * o * sqrt_det_cov_g * sqrt_det_cov * c_without_2pi(mu_g, mu, cov_g, cov)
+
+        return (f2integral + g2integral + h2integral + 2 * fgintegral - 2 * fhintegral - 2 * ghintegral) * (np.pi ** 1.5)
+
+
+    def target_inner(features) -> float:
+        mu = features[:3]
+        s = features[3:6]
+        r = features[6:10]
+        o = features[10]
+        gaussian_h = make_gaussian(mu, s, r, o)
+
+        N = 64
+
+        sum = 0.
+        tot_weight = 0.
+        random.seed(int(time.time()))
+        for _ in range(N):
+            while True:
+                x = random.random()
+                y = random.random()
+                z = random.random()
+                pos = point_min + x * basis_x + y * basis_y + z * basis_z
+
+                f_value = gaussian_f(pos)
+                g_value = gaussian_g(pos)
+
+                if N < 32: # 用 f 的 pdf
+                    weight = f_value / opacity_f
+                    if random.random() < weight:
+                        break
+                else: # 用 g 的 pdf
+                    weight = g_value / opacity_g
+                    if random.random() < weight:
+                        break
+
+            sum += (f_value + g_value - gaussian_h(pos)) ** 2 * weight
+            tot_weight += weight
+        return sum / tot_weight
+    
+    res = minimize(target_inner_analytical, x_0, options={'gtol': 1e-4, 'disp': False})
+    mu_h = res.x[:3]
+    s_h = res.x[3:6]
+    r_h = res.x[6:10]
+    o_h = res.x[10]
+
+    sh = (gaussian.sh[id1] + gaussian.sh[id2]) / 2.
+    features_dc = sh[0, :]
+    features_rest = sh[1:, :]
+    # print(mu_h, s_h, r_h, features_dc, features_rest, o_h)
+    
+    return mu_h, s_h, r_h, features_dc, features_rest, o_h
+
+# def merge_neo(gaussian: Gaussian, id1: int, id2: int):
+#     N_f = np.prod(gaussian.scales[id1]) * gaussian.opacity[id1]
+#     N_g = np.prod(gaussian.scales[id2]) * gaussian.opacity[id2]
+#     # N_f = gaussian.opacity[id1]
+#     # N_g = gaussian.opacity[id2]
+
+#     mu_f = gaussian.positions[id1]
+#     mu_g = gaussian.positions[id2]
+#     Lambda_f = np.diag(gaussian.scales[id1] ** 2)
+#     Lambda_g = np.diag(gaussian.scales[id2] ** 2)
+#     U_f = quat2rot(gaussian.rotations[id1])
+#     U_g = quat2rot(gaussian.rotations[id2])
+
+#     N_h = N_f + N_g
+#     position = (N_f * mu_f + N_g * mu_g) / N_h
+    
+#     g = U_f.T @ (mu_f - mu_g)
+#     G = U_f.T @ U_g
+
+#     mat = N_f / N_h * Lambda_f + \
+#         N_g / N_h * G @ Lambda_g @ G.T + \
+#         N_f * N_g / (N_h * N_h) * g @ g.T
+#     eigenvalues, eigenvectors = np.linalg.eig(mat)
+#     U_h = U_f @ eigenvectors
+
+#     rotation = rot2quat(U_h)
+#     scale = np.sqrt(eigenvalues)
+#     opacity = N_h / np.prod(scale)
+
+#     sh = (gaussian.sh[id1] + gaussian.sh[id2]) / 2.
+#     features_dc = sh[0, :]
+#     features_rest = sh[1:, :]
+
+#     return position, scale, rotation, features_dc, features_rest, opacity
+
+def merge_geo(gaussian: Gaussian, id1: int, id2: int):
+    rotation = gaussian.rotations[id1] + gaussian.rotations[id2]
+    rotation = rotation / np.linalg.norm(rotation)
+
+    rotation_mat = quat2rot(rotation)
+    inv_rotation_mat = np.linalg.inv(rotation_mat)
+
+    rotation1_mat = inv_rotation_mat @ quat2rot(gaussian.rotations[id1])
+    rotation2_mat = inv_rotation_mat @ quat2rot(gaussian.rotations[id2])
+    position1 = inv_rotation_mat @ gaussian.positions[id1]
+    position2 = inv_rotation_mat @ gaussian.positions[id2]
+    aabb1 = process_aabb(position1, gaussian.scales[id1], rotation1_mat)
+    aabb2 = process_aabb(position2, gaussian.scales[id2], rotation2_mat)
+    aabb = aabb_merge(aabb1, aabb2)
+
+    # point_min = aabb[:3]
+    # vectors = aabb[3:] - point_min
+    # vector_x = rotation_mat @ np.array([vectors.x, 0, 0])
+    # vector_y = rotation_mat @ np.array([0, vectors.y, 0])
+    # vector_z = rotation_mat @ np.array([0, 0, vectors.z])
+    
+    point_min = np.array([aabb[0], aabb[2], aabb[4]])
+    point_max = np.array([aabb[1], aabb[3], aabb[5]])
+    vectors = rotation_mat @ np.diag(point_max - point_min)
+
+    # new_aabb = np.concatenate((rotation_mat @ point_min, rotation_mat @ point_max), axis=0)
+
+    position = rotation_mat @ center(aabb)
+    scale = np.array([aabb[1] - aabb[0], aabb[3] - aabb[2], aabb[5] - aabb[4]]) / 2.
+    opacity = np.array((gaussian.opacity[id1] + gaussian.opacity[id2]) / 2.)
+    return position, scale, rotation, opacity, rotation_mat @ point_min, vectors
+
+
+def merge_2(gaussian: Gaussian, id1: int, id2: int):
+    rotation = gaussian.rotations[id1] + gaussian.rotations[id2]
+    rotation = rotation / np.linalg.norm(rotation)
+
+    rotation_mat = quat2rot(rotation)
+    inv_rotation_mat = np.linalg.inv(rotation_mat)
+
+    rotation1_mat = inv_rotation_mat @ quat2rot(gaussian.rotations[id1])
+    rotation2_mat = inv_rotation_mat @ quat2rot(gaussian.rotations[id2])
+    position1 = inv_rotation_mat @ gaussian.positions[id1]
+    position2 = inv_rotation_mat @ gaussian.positions[id2]
+    aabb1 = process_aabb(position1, gaussian.scales[id1], rotation1_mat)
+    aabb2 = process_aabb(position2, gaussian.scales[id2], rotation2_mat)
+    aabb = aabb_merge(aabb1, aabb2)
+    
+    position = rotation_mat @ center(aabb)
+    scale = np.array([aabb[1] - aabb[0], aabb[3] - aabb[2], aabb[5] - aabb[4]]) / 2.
+
+    sh = (gaussian.sh[id1] + gaussian.sh[id2]) / 2.
+    features_dc = sh[0, :]
+    features_rest = sh[1:, :]
+    opacity = (gaussian.opacity[id1] + gaussian.opacity[id2]) / 2.
+    return position, scale, rotation, features_dc, features_rest, opacity
+
+def merge(gaussian: Gaussian, id1: int, id2: int):
+    aabb1 = process_aabb(
+        gaussian.positions[id1],
+        gaussian.scales[id1],
+        gaussian.rotations[id1])
+    aabb2 = process_aabb(
+        gaussian.positions[id2],
+        gaussian.scales[id2],
+        gaussian.rotations[id2])
+    aabb = aabb_merge(aabb1, aabb2)
+    print(f"merging:\n  {aabb1}\n  {aabb2}\n\n  {aabb}")
+
+    position = center(aabb)
+    rotation = gaussian.rotations[id1] + gaussian.rotations[id2]
+    rotation = rotation / np.linalg.norm(rotation)
+    scale = solve_scale(rotation, aabb)
+    if np.isnan(scale).any():
+        # 为什么会出现 nan??????
+        return None, None, None, None, None, None
+    sh = (gaussian.sh[id1] + gaussian.sh[id2]) / 2.
+    features_dc = sh[0, :]
+    features_rest = sh[1:, :]
+    opacity = (gaussian.opacity[id1] + gaussian.opacity[id2]) / 2.
+    # opacity = min(gaussian.opacity[id1], gaussian.opacity[id2])
+
+    # aabb_test = process_aabb(
+    #     position,
+    #     scale,
+    #     rotation,
+    #     scale_factor=1
+    # )
+    # assert(np.allclose(aabb, aabb_test))
+
+    return position, scale, rotation, features_dc, features_rest, opacity
+    # gaussian.replace(id1, id2, position, scale, rotation, sh, opacity)
diff --git a/new_main.py b/new_main.py
new file mode 100644
index 0000000..99a6767
--- /dev/null
+++ b/new_main.py
@@ -0,0 +1,129 @@
+from dataloader import load_ply, save_ply
+from gaussian import Gaussian
+from scipy.spatial import KDTree
+from scipy.optimize import minimize
+from merge import merge_geo
+from utils import quat2rot
+from tqdm import tqdm
+import argparse
+import numpy as np
+import pickle as pkl
+
+NUMBER=8000
+
+parser = argparse.ArgumentParser("main.py")
+parser.add_argument('-m', '--model_path', type=str, default="/home/chuyan/Documents/code/models/train/point_cloud/iteration_50000/point_cloud.ply")
+parser.add_argument('-o', '--output_path', type=str, default="/home/chuyan/Documents/code/models/train/point_cloud/iteration_60000/point_cloud.ply")
+parser.add_argument('-d', '--sh_degrees', type=int, default=3)
+
+args = parser.parse_args()
+
+gaussian_model = load_ply(args.model_path, args.sh_degrees)
+gaussian_extended = Gaussian.empty_with_cap(NUMBER)
+
+kd_tree = KDTree(gaussian_model.positions)
+pairs = []
+for i in tqdm(range(gaussian_model.num_gaussians)):
+    scale = gaussian_model.scales[i]
+    radius = max(max(scale[0], scale[1]), scale[2]) * 1.5
+    points = kd_tree.query_ball_point(gaussian_model.positions[i], r=radius, workers=18)
+    for j in points:
+        if j < i:
+            pairs.append((i, j))
+
+# with open('pairs.pkl', 'wb') as f:
+#     pkl.dump(pairs, f)
+
+def merge_3(gaussian: Gaussian, id1: int, id2: int) -> float:
+    mu_f = gaussian.positions[id1]
+    mu_g = gaussian.positions[id2]
+    s_f = gaussian.scales[id1]
+    s_g = gaussian.scales[id2]
+    r_f = gaussian.rotations[id1]
+    r_g = gaussian.rotations[id2]
+    o_f = gaussian.opacity[id1]
+    o_g = gaussian.opacity[id2]
+
+    mu_0, s_0, r_0, o_0, point_min, vectors = merge_geo(gaussian_model, id1, id2)
+    vector_x = vectors[:, 0]
+    vector_y = vectors[:, 1]
+    vector_z = vectors[:, 2]
+    x_0 = np.concatenate((mu_0, s_0, r_0, o_0), axis=0)
+
+    def make_gaussian(mu, s, r, o):
+        def calc_inner(pos):
+            pos = pos - mu
+            S = np.diag(s)
+            R = quat2rot(r)
+            M = S @ R
+            Sigma = np.linalg.inv(M.T @ M)
+            return o * np.exp(-0.5 * np.dot(pos, Sigma @ pos))
+        return calc_inner
+    gaussian_f = make_gaussian(mu_f, s_f, r_f, o_f)
+    gaussian_g = make_gaussian(mu_g, s_g, r_g, o_g)
+
+    def target_inner(features) -> float:
+        mu = features[:3]
+        s = features[3:6]
+        r = features[6:10]
+        o = features[10:]
+        gaussian_h = make_gaussian(mu, s, r, o)
+
+        N = 4
+
+        sum = 0.
+        for i in range(N):
+            for j in range(N):
+                for k in range(N):
+                    pos = point_min + (i * vector_x + j * vector_y + k * vector_z) / N
+                    sum += (gaussian_f(pos) + gaussian_g(pos) - gaussian_h(pos)) ** 2
+        return sum
+    
+    res = minimize(target_inner, x_0)
+    mu_h = res.x[:3]
+    s_h = res.x[3:6]
+    r_h = res.x[6:10]
+    o_h = res.x[10]
+
+    sh = (gaussian.sh[id1] + gaussian.sh[id2]) / 2.
+    features_dc = sh[0, :]
+    features_rest = sh[1:, :]
+    # print(mu_h, s_h, r_h, features_dc, features_rest, o_h)
+    
+    return mu_h, s_h, r_h, features_dc, features_rest, o_h
+
+cnt = 0
+filter = np.ones(gaussian_model.num_gaussians, dtype=bool)
+used = set()
+for i, j in tqdm(pairs):
+    if i in used or j in used:
+        continue
+    mu, s, r, dc, rest, o = merge_3(gaussian_model, i, j)
+    if mu is not None:
+        # scale_product = np.prod(scale)
+        # scale_product_1 = np.prod(gaussian_model.scales[i])
+        # scale_product_2 = np.prod(gaussian_model.scales[j])
+        # if scale_product / max(scale_product_1, scale_product_2) > 10.:
+        #     continue
+        
+        cnt += 1
+        id = gaussian_extended.add(mu, s, r, dc, rest, o)
+        filter[i] = False
+        filter[j] = False
+        used.add(i)
+        used.add(j)
+
+        if cnt == NUMBER:
+            break
+
+deleted_gaussian = gaussian_model.copy()
+deleted_gaussian.apply_filter(np.logical_not(filter))
+
+gaussian_model.apply_filter(filter)
+gaussian_model.concat(gaussian_extended)
+
+deleted_path = args.output_path.replace('60000', '60001')
+added_path = args.output_path.replace('60000', '60002')
+save_ply(args.output_path, gaussian_model)
+save_ply(deleted_path, deleted_gaussian)
+save_ply(added_path, gaussian_extended)
\ No newline at end of file
diff --git a/similarity.py b/similarity.py
new file mode 100644
index 0000000..fec83d6
--- /dev/null
+++ b/similarity.py
@@ -0,0 +1,69 @@
+from gaussian import Gaussian
+from utils import quat2rot
+from aabb import process_aabb, intersect, min_interval
+import numpy as np
+
+def similarity(gaussian: Gaussian, id1: int, id2: int):
+    geometry_sim = geometry_similarity(gaussian, id1, id2)
+    color_sim = color_similarity(gaussian.sh[id1], gaussian.sh[id2])
+    color_max = np.max(np.abs(gaussian.sh))
+    # return geometry_sim * geometry_coeff + color_sim * color_coeff
+    return geometry_sim, color_sim / (color_max * color_max * 48)
+    
+def color_similarity(c1: np.ndarray, c2: np.ndarray):
+    return np.sum(c1 * c2)
+
+def geometry_similarity(gaussian: Gaussian, id1: int, id2: int):
+    def geometry_similarity_using_overlap(gaussian: Gaussian, id1: int, id2: int):
+        aabb1 = process_aabb(gaussian.positions[id1], gaussian.scales[id1], gaussian.rotations[id1], scale_factor=2)
+        aabb2 = process_aabb(gaussian.positions[id2], gaussian.scales[id2], gaussian.rotations[id2], scale_factor=2)
+        # 在两个 aabb 里 stratified sample 点，对两个 Gaussian 的乘积做积分
+        # 然后除以总sample数作为重叠程度的 metric
+
+        intersected_aabb = intersect(aabb1, aabb2)
+
+        if intersected_aabb is None:
+            return 0
+
+        x_stride = (intersected_aabb[1] - intersected_aabb[0]) / (4. + 1e-3)
+        y_stride = (intersected_aabb[3] - intersected_aabb[2]) / (4. + 1e-3)
+        z_stride = (intersected_aabb[5] - intersected_aabb[4]) / (4. + 1e-3)
+        
+        score = 0
+        for i in range(int(np.floor((intersected_aabb[1] - intersected_aabb[0]) / x_stride))):
+            x = intersected_aabb[0] + i * x_stride
+            for j in range(int(np.floor((intersected_aabb[3] - intersected_aabb[2]) / y_stride))):
+                y = intersected_aabb[2] + j * y_stride
+                for k in range(int(np.floor((intersected_aabb[5] - intersected_aabb[4]) / z_stride))):
+                    z = intersected_aabb[4] + k * z_stride
+                    pos = np.array([x, y, z])
+                    score += gaussian.calc(pos, id1, 10) * gaussian.calc(pos, id2, 10) * gaussian.opacity[id1] * gaussian.opacity[id2]
+
+        scale1 = gaussian.scales[id1][0] * gaussian.scales[id1][1] * gaussian.scales[id1][2]
+        scale2 = gaussian.scales[id2][0] * gaussian.scales[id2][1] * gaussian.scales[id2][2]
+
+        # print(f"calculating pair ({id1},{id2}), used {64} samples.")
+        return score / (scale1 * scale2)
+
+    def geometry_similarity_using_position(gaussian: Gaussian, id1: int, id2: int):
+        def position_similarity(p1: np.ndarray, p2: np.ndarray):
+            return np.linalg.norm(p1 - p2)
+        
+        def scale_similarity(s1: np.ndarray, s2: np.ndarray):
+            return np.linalg.norm(s1 - s2)
+
+        def rotation_similarity(q1: np.ndarray, q2: np.ndarray):
+            return 1 - np.dot(q1, q2)
+            # r1 = quat2rot(q1)
+            # r2 = quat2rot(q2)
+            # eigenvalues, eigenvectors = np.linalg.eig(r1 @ r2.T)
+            # return np.linalg.norm(eigenvectors @ np.diag(np.log(eigenvalues)) @ np.linalg.inv(eigenvectors)) 
+        position_coeff = 1.
+        scale_coeff = 1.
+        rotation_coeff = 1.
+        return position_similarity(gaussian.positions[id1], gaussian.positions[id2]) * position_coeff + \
+            scale_similarity(gaussian.scales[id1], gaussian.scales[id2]) * scale_coeff + \
+            rotation_similarity(gaussian.rotations[id1], gaussian.rotations[id2]) * rotation_coeff
+
+    return geometry_similarity_using_overlap(gaussian, id1, id2)
+
diff --git a/utils.py b/utils.py
new file mode 100644
index 0000000..95a1c31
--- /dev/null
+++ b/utils.py
@@ -0,0 +1,36 @@
+import numpy as np
+
+def quat2rot(q):
+    r, x, y, z = q[0], q[1], q[2], q[3]
+    return np.array([
+        [1. - 2. * (y * y + z * z), 2. * (x * y - r * z), 2. * (x * z + r * y)],
+		[2. * (x * y + r * z), 1. - 2. * (x * x + z * z), 2. * (y * z - r * x)],
+		[2. * (x * z - r * y), 2. * (y * z + r * x), 1. - 2. * (x * x + y * y)]
+    ], dtype=np.float32)
+
+def rot2quat(r):
+    if r[2][2] < 0:
+        if r[0][0] > r[1][1]:
+            t = 1 + r[0][0] - r[1][1] - r[2][2]
+            q = np.array([t, r[0][1] + r[1][0], r[0][2] + r[2][0], r[1][2] - r[2][1]])
+        else:
+            t = 1 - r[0][0] + r[1][1] - r[2][2]
+            q = np.array([r[0][1] + r[1][0], t, r[1][2] + r[2][1], r[2][0] - r[0][2]])
+    else:
+        if r[0][0] < -r[1][1]:
+            t = 1 - r[0][0] - r[1][1] + r[2][2]
+            q = np.array([r[2][0] + r[0][2], r[1][2] + r[2][1], t, r[0][1] - r[1][0]])
+        else:
+            t = 1 + r[0][0] + r[1][1] + r[2][2]
+            q = np.array([r[1][2] - r[2][1], r[2][0] + r[0][2], r[0][1] - r[1][0], t])
+    return q * 0.5 / np.sqrt(t)
+
+def sq2cov(s, q):
+    S = np.diag(s)
+    R = quat2rot(q)
+    M = S @ R
+    return M.T @ M
+
+def mat_log(m):
+    eigenvalues, eigenvectors = np.linalg.eig(m)
+    return eigenvectors @ np.diag(np.log(eigenvalues)) @ np.linalg.inv(eigenvectors)
\ No newline at end of file