From 394e19b012cb9264feaec582948fa7ac8bff901c Mon Sep 17 00:00:00 2001 From: Chuyan Zhang Date: Mon, 15 Jan 2024 15:49:41 -0800 Subject: init commit --- merge.py | 259 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 259 insertions(+) create mode 100644 merge.py (limited to 'merge.py') 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) -- cgit v1.2.3-70-g09d2