From 394e19b012cb9264feaec582948fa7ac8bff901c Mon Sep 17 00:00:00 2001 From: Chuyan Zhang Date: Mon, 15 Jan 2024 15:49:41 -0800 Subject: init commit --- .gitignore | 3 + aabb.py | 102 +++++++++++++++++++++++ dataloader.py | 116 ++++++++++++++++++++++++++ gaussian.py | 116 ++++++++++++++++++++++++++ main.py | 99 ++++++++++++++++++++++ merge.py | 259 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ new_main.py | 129 +++++++++++++++++++++++++++++ similarity.py | 69 ++++++++++++++++ utils.py | 36 ++++++++ 9 files changed, 929 insertions(+) create mode 100644 .gitignore create mode 100644 aabb.py create mode 100644 dataloader.py create mode 100644 gaussian.py create mode 100644 main.py create mode 100644 merge.py create mode 100644 new_main.py create mode 100644 similarity.py create mode 100644 utils.py 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 -- cgit v1.2.3-70-g09d2