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)