import os import sys import random import torch import numpy as np import torch.nn.functional as F from argparse import ArgumentParser from core.registry import register_method from core.base_method import BaseMethod sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../GeoGaussian'))) from utils.loss_utils import l1_loss, ssim from gaussian_renderer import render as native_render from diff_gaussian_rasterization import SurfaceAlign from scene import Scene, GaussianModel from arguments import ModelParams, PipelineParams, OptimizationParams @register_method("geogaussian") class GeoGaussianWrapper(BaseMethod): def __init__(self, dataset_config, hyperparams): self.parser = ArgumentParser() self.lp = ModelParams(self.parser) self.op = OptimizationParams(self.parser) self.pp = PipelineParams(self.parser) self.args = self.parser.parse_args([]) self.args.source_path = dataset_config["source_path"] self.args.model_path = dataset_config["model_path"] self.args.eval = True self.args.resolution = dataset_config.get("resolution", 1) self.track_decoupling = hyperparams.get("track_decoupling", False) self.dataset = self.lp.extract(self.args) self.opt = self.op.extract(self.args) self.pipe = self.pp.extract(self.args) self.gaussians = GaussianModel(self.dataset.sh_degree) # INJECTED_RES_FIX begin import sys as _sys _scene, _explicit_res = None, None for _i, _a in enumerate(_sys.argv[:-1]): _v = _sys.argv[_i + 1] if _a == "--scene": _scene = _v elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] elif _a == "--resolution": try: _explicit_res = int(_v) except: pass _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} if _explicit_res is not None and _explicit_res > 0: _res = _explicit_res elif _scene is not None: _res = 4 if _scene in _OUTDOOR_360 else 2 else: _res = None try: if _res is not None: self.dataset.resolution = _res print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) except Exception as _e: print("[res-fix] FAILED:", _e) # INJECTED_RES_FIX end self.scene = Scene(self.dataset, self.gaussians) self.gaussians.training_setup(self.opt) bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") self.viewpoint_stack = self.scene.getTrainCameras().copy() self.last_n_gaussians = len(self.gaussians.get_xyz) self.KNN_index = None def train_iteration(self, step): self.gaussians.update_learning_rate(step) if step % 1000 == 0: self.gaussians.oneupSHdegree() if not self.viewpoint_stack: self.viewpoint_stack = self.scene.getTrainCameras().copy() viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) bg = torch.rand((3), device="cuda") if self.opt.random_background else self.background render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, bg) image = render_pkg["render"] viewspace_point_tensor = render_pkg["viewspace_points"] visibility_filter = render_pkg["visibility_filter"] radii = render_pkg["radii"] gt_image = viewpoint_cam.original_image.cuda() Ll1 = l1_loss(image, gt_image) ssim_value = ssim(image, gt_image) loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 loss_parasitic_ssim = self.opt.lambda_dssim * (1.0 - ssim_value) pair_d_loss = torch.tensor(0.0, device="cuda") pair_normal_loss = torch.tensor(0.0, device="cuda") if self.KNN_index is not None and visibility_filter.sum() > 0: mask = torch.logical_and(visibility_filter, self.gaussians.get_type.squeeze() == 1) if mask.sum() > 0: pair_d_loss_raw, pair_normal_loss_raw = SurfaceAlign()( self.gaussians.get_xyz, self.gaussians.get_xyz_id.contiguous(), self.gaussians.get_rotation, self.KNN_index[mask] ) pair_d_loss = 0.05 * torch.mean(pair_d_loss_raw) pair_normal_loss = 0.01 * torch.mean(pair_normal_loss_raw) loss_parasitic = loss_parasitic_ssim + pair_d_loss + pair_normal_loss loss = loss_target + loss_parasitic grad_cos_sim = 0.0 parasitic_ratio = 0.0 stats = {} if self.track_decoupling and step % 100 == 0: self.gaussians.optimizer.zero_grad(set_to_none=True) loss_target.backward(retain_graph=True) grads_target = {} for group in self.gaussians.optimizer.param_groups: name = group["name"] p = group["params"][0] if p.grad is not None: grads_target[name] = p.grad.clone() self.gaussians.optimizer.zero_grad(set_to_none=True) loss_parasitic.backward(retain_graph=True) grads_parasitic = {} for group in self.gaussians.optimizer.param_groups: name = group["name"] p = group["params"][0] if p.grad is not None: grads_parasitic[name] = p.grad.clone() gt_xyz = grads_target.get("xyz", torch.zeros(1, device="cuda")) gp_xyz = grads_parasitic.get("xyz", torch.zeros(1, device="cuda")) valid_mask = (torch.norm(gt_xyz, dim=1) > 0) & (torch.norm(gp_xyz, dim=1) > 0) if valid_mask.any(): grad_cos_sim = float(F.cosine_similarity(gt_xyz[valid_mask], gp_xyz[valid_mask], dim=1).mean()) parasitic_ratio = float(torch.norm(gp_xyz, dim=1).mean() / (torch.norm(gt_xyz, dim=1).mean() + 1e-7)) group_mapping = { "spatial": ["xyz"], "geometry": ["scaling", "rotation"], "opacity": ["opacity"], "appearance": ["f_dc", "f_rest"] } for super_group, sub_names in group_mapping.items(): u_t_list = [] u_p_list = [] for name in sub_names: for group in self.gaussians.optimizer.param_groups: if group["name"] == name: p = group["params"][0] state = self.gaussians.optimizer.state.get(p) vt = state.get("exp_avg_sq", torch.zeros_like(p)) if state else torch.zeros_like(p) lr = group["lr"] gt = grads_target.get(name, torch.zeros_like(p)) gp = grads_parasitic.get(name, torch.zeros_like(p)) step_t = (lr / (torch.sqrt(vt) + 1e-15)) * gt step_p = (lr / (torch.sqrt(vt) + 1e-15)) * gp u_t_list.append(step_t.reshape(-1)) u_p_list.append(step_p.reshape(-1)) if u_t_list and u_p_list: U_T = torch.cat(u_t_list) U_P = torch.cat(u_p_list) if U_T.norm() > 0 and U_P.norm() > 0: cos = float(F.cosine_similarity(U_T.unsqueeze(0), U_P.unsqueeze(0))) r = float(U_P.norm() / (U_T.norm() + U_P.norm() + 1e-7)) ti = r * max(0.0, -cos) else: ti = 0.0 stats[f"sti_{super_group}"] = ti else: stats[f"sti_{super_group}"] = 0.0 self.gaussians.optimizer.zero_grad(set_to_none=True) loss.backward() else: loss.backward() with torch.no_grad(): if step < self.opt.densify_until_iter: self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: size_threshold = 20 if step > self.opt.opacity_reset_interval else None self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.05, self.scene.cameras_extent, size_threshold) self.KNN_index = self.gaussians.findKNN() if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): self.gaussians.reset_opacity() elif step > self.opt.densify_until_iter: if step % 3000 == 0: self.KNN_index = self.gaussians.findKNN() if step <= self.opt.iterations: self.gaussians.optimizer.step() self.gaussians.optimizer.zero_grad(set_to_none=True) num_gaussians = self.gaussians.get_xyz.shape[0] metrics = { "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic_ssim), "loss_pair_d": float(pair_d_loss), "loss_pair_normal": float(pair_normal_loss), "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio) } metrics.update(stats) self.last_n_gaussians = num_gaussians histograms = {} if step % 1000 == 0: histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() scales = torch.exp(self.gaussians._scaling).clone().detach() histograms["scaling"] = scales scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) histograms["anisotropy"] = gamma histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) return metrics, histograms def render(self, camera): with torch.no_grad(): render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) return {"image": render_pkg["render"], "depth": render_pkg.get("render_depth", None)} def save(self, save_dir, step): self.scene.save(step) def load(self, model_path, iteration): self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) def get_spatial_centers(self): return self.gaussians._xyz def compute_physical_metrics(self, cameras=None): metrics = {} with torch.no_grad(): raw_scales = self.gaussians._scaling scales = torch.exp(raw_scales) scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) max_S, _ = torch.max(scales_2d, dim=1) min_S, _ = torch.min(scales_2d, dim=1) gamma = max_S / (min_S + 1e-7) metrics["gamma_median"] = float(torch.median(gamma)) metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) metrics["scale_mean"] = float(torch.mean(scales_2d)) metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) dc, rest = self.gaussians._features_dc, self.gaussians._features_rest if rest is not None and rest.shape[1] > 0: metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) if cameras is not None and len(cameras) > 0: view_dirs = [] for c in cameras: view_dirs.append(c.world_view_transform[:3, 2].tolist()) view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) rots = F.normalize(self.gaussians._rotation.clone(), dim=1) w, x, y, z = rots.unbind(dim=-1) normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) return metrics def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: with torch.no_grad(): V = query_points.shape[0] densities = torch.zeros(V, device="cuda") xyz = self.gaussians._xyz opacities = torch.sigmoid(self.gaussians._opacity).squeeze() scales = torch.exp(self.gaussians._scaling) sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) N_gaussians = xyz.shape[0] chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) for i in range(0, V, chunk_size): end = min(i + chunk_size, V) dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) return densities