methods/wrapper_gspull.py

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")