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 GS_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '../../GS-Pull')) # 强隔离：清空残留路径并确保 GS-Pull 绝对优先 if GS_ROOT not in sys.path: sys.path.insert(0, GS_ROOT) # GS-Pull's gaussian_splatting is a vanilla 3DGS copy — its scene/__init__.py # does `from utils.system_utils import ...` expecting the gaussian_splatting/ # subdir to be on sys.path. _GS_INNER = os.path.join(GS_ROOT, 'gaussian_splatting') if _GS_INNER not in sys.path: sys.path.insert(0, _GS_INNER) from gaussian_splatting.scene import Scene, GaussianModel from gaussian_splatting.arguments import ModelParams, PipelineParams, OptimizationParams as GSOptParams from gaussian_splatting.gaussian_renderer import render as vanilla_render from sugar_scene.sugar_model import SuGaR from sugar_scene.sugar_optimizer import OptimizationParams as SuGaROptParams, SuGaROptimizer from sugar_scene.sugar_densifier import SuGaRDensifier from sugar_utils.loss_utils import ssim, l1_loss from np_utils.train import Runner @register_method("gspull") class GSPullWrapper(BaseMethod): def __init__(self, dataset_config, hyperparams): self.source_path = dataset_config["source_path"] self.model_path = dataset_config["model_path"] self.resolution = dataset_config.get("resolution", 1) self.track_decoupling = hyperparams.get("track_decoupling", False) parser = ArgumentParser() lp = ModelParams(parser) op = GSOptParams(parser) pp = PipelineParams(parser) args = parser.parse_args([]) args.source_path = self.source_path args.model_path = self.model_path args.resolution = self.resolution args.eval = True self.dataset = lp.extract(args) self.gs_opt = op.extract(args) self.pipe = pp.extract(args) # 防呆逻辑：兼容某些变体改名为 shfeature_lr 的情况 if not hasattr(self.gs_opt, 'shfeature_lr') and hasattr(self.gs_opt, 'feature_lr'): setattr(self.gs_opt, 'shfeature_lr', getattr(self.gs_opt, 'feature_lr')) 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.gs_opt) self.bg_color = torch.tensor([1, 1, 1] if self.dataset.white_background else [0, 0, 0], dtype=torch.float32, device="cuda") self.train_cameras = self.scene.getTrainCameras().copy() self.viewpoint_stack = self.train_cameras.copy() self.sugar = None self.optimizer = None self.densifier = None self._iter = 0 self.last_n_gaussians = len(self.gaussians.get_xyz) def _init_sugar_stage(self): self.sugar = SuGaR( nerfmodel=None, points=self.gaussians.get_xyz.detach(), colors=torch.zeros_like(self.gaussians.get_xyz), initialize=False, sh_levels=self.gaussians.active_sh_degree + 1, learnable_positions=True, triangle_scale=1.0, keep_track_of_knn=True, knn_to_track=16, beta_mode="average", freeze_gaussians=False, surface_mesh_to_bind=None, ) with torch.no_grad(): self.sugar._points.copy_(self.gaussians._xyz) self.sugar._scales.copy_(self.gaussians._scaling) self.sugar._quaternions.copy_(self.gaussians._rotation) self.sugar.all_densities.copy_(self.gaussians._opacity) self.sugar._sh_coordinates_dc.copy_(self.gaussians._features_dc) self.sugar._sh_coordinates_rest.copy_(self.gaussians._features_rest) self.sugar.part_num = 1 self.sugar.neus = Runner(self.model_path, None, part_num=self.sugar.part_num) spatial_lr_scale = self.scene.cameras_extent opt_p = SuGaROptParams( iterations=15000, position_lr_init=0.00016, position_lr_final=0.0000016, position_lr_max_steps=30000, scaling_lr=0.005, opacity_lr=0.05, feature_lr=0.0025, rotation_lr=0.001 ) self.optimizer = SuGaROptimizer(self.sugar, opt_p, spatial_lr_scale=spatial_lr_scale) self.densifier = SuGaRDensifier(self.sugar, self.optimizer, 0.0002, 0.005, 20, spatial_lr_scale, 0.01) def train_iteration(self, step): self._iter += 1 metrics = {} histograms = {} if self._iter <= 7000: if not self.viewpoint_stack: self.viewpoint_stack = self.train_cameras.copy() viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) render_pkg = vanilla_render(viewpoint_cam, self.gaussians, self.pipe, self.bg_color) image, viewspace_point_tensor, visibility_filter, radii = render_pkg["render"], render_pkg["viewspace_points"], render_pkg["visibility_filter"], render_pkg["radii"] gt_image = viewpoint_cam.original_image.cuda() loss = 0.8 * l1_loss(image, gt_image) + 0.2 * (1.0 - ssim(image, gt_image)) loss.backward() with torch.no_grad(): 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 self._iter > 500 and self._iter % 100 == 0: self.gaussians.densify_and_prune(self.gs_opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, 20) if self._iter % 3000 == 0: self.gaussians.reset_opacity() self.gaussians.optimizer.step() self.gaussians.optimizer.zero_grad(set_to_none=True) num_gaussians = len(self.gaussians.get_xyz) metrics = { "loss": float(loss), "num_gaussians": int(num_gaussians), "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)) } self.last_n_gaussians = num_gaussians return metrics, histograms if self._iter == 7001: self._init_sugar_stage() is_pulling = self._iter > 9000 if self._iter == 9001: prune_mask = (self.sugar.strengths < 0.5).squeeze() self.densifier.prune_points(prune_mask) self.sugar.reset_neighbors() self.sugar.neus.reset_datasets(self.model_path, self.sugar.points.detach().cpu().numpy(), iteration=9000, scene_name="scene") if not self.viewpoint_stack: self.viewpoint_stack = self.train_cameras.copy() cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) outputs = self.sugar.render_image_gaussian_rasterizer( camera_indices=self.train_cameras.index(cam), bg_color=self.bg_color, sh_deg=self.sugar.sh_levels-1, compute_covariance_in_rasterizer=True, return_2d_radii=True ) pred_rgb = outputs['image'].view(-1, cam.image_height, cam.image_width, 3).transpose(-1, -2).transpose(-2, -3) gt_rgb = cam.original_image.cuda().unsqueeze(0) loss = 0.8 * l1_loss(pred_rgb, gt_rgb) + 0.2 * (1.0 - ssim(pred_rgb, gt_rgb)) loss.backward() self.optimizer.step() self.optimizer.zero_grad(set_to_none=True) metrics = {"loss": float(loss), "num_gaussians": len(self.sugar.points)} return metrics, {} def render(self, camera): with torch.no_grad(): if self._iter <= 7000: return {"image": vanilla_render(camera, self.gaussians, self.pipe, self.bg_color)["render"], "depth": None} idx = 0 for i, c in enumerate(self.train_cameras): if c.uid == camera.uid: idx = i; break outputs = self.sugar.render_image_gaussian_rasterizer(camera_indices=idx) return {"image": outputs["image"], "depth": None} def save(self, save_dir, step): if self.sugar: self.sugar.save_model(path=os.path.join(save_dir, f"{step}.pt"), iteration=step) else: self.gaussians.save_ply(os.path.join(save_dir, "point_cloud", f"iteration_{step}", "point_cloud.ply")) def load(self, model_path, iteration): pass def get_spatial_centers(self): return self.sugar.points if self.sugar else self.gaussians.get_xyz def compute_physical_metrics(self, cameras=None): return {"num_gaussians": float(len(self.get_spatial_centers()))} def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: return torch.zeros(len(query_points), device="cuda")