summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--.gitignore3
-rw-r--r--aabb.py102
-rw-r--r--dataloader.py116
-rw-r--r--gaussian.py116
-rw-r--r--main.py99
-rw-r--r--merge.py259
-rw-r--r--new_main.py129
-rw-r--r--similarity.py69
-rw-r--r--utils.py36
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