| import os |
| import sys |
| import torch |
| import random |
| 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 |
|
|
| SURFEL_ROOT = "/root/autodl-tmp/gaussian_surfels" |
| if SURFEL_ROOT not in sys.path: |
| sys.path.insert(0, SURFEL_ROOT) |
|
|
| @register_method("gaussian_surfel") |
| class SurfelWrapper(BaseMethod): |
| def __init__(self, dataset_config, hyperparams): |
| from scene import Scene, GaussianModel |
| from arguments import ModelParams, PipelineParams, OptimizationParams |
| 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.resolution = dataset_config.get("resolution", 1) |
| self.args.eval = True |
| self.args.surface = True |
| if any(x in self.args.source_path.lower() for x in ["tnt", "360", "tanks"]): |
| self.args.images = "images" |
| else: |
| self.args.images = "image" |
| self.dataset = self.lp.extract(self.args) |
| self.opt = self.op.extract(self.args) |
| self.pipe = self.pp.extract(self.args) |
| _prev_cwd = os.getcwd(); os.chdir("/root/autodl-tmp/gaussian_surfels") |
| self.gaussians = GaussianModel(self.dataset) |
| os.chdir(_prev_cwd) |
| |
| 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) |
| |
| self.scene = Scene(self.dataset, self.gaussians, self.opt.camera_lr, shuffle=False, resolution_scales=[1, 2, 4]) |
| 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 = None |
| self.last_n_gaussians = len(self.gaussians.get_xyz) |
| self.track_decoupling = hyperparams.get("track_decoupling", False) |
| self.cap_gaussians = hyperparams.get("cap_gaussians", None) |
|
|
| def train_iteration(self, step): |
| from gaussian_renderer import render as native_render |
| from utils.loss_utils import l1_loss, ssim, cos_loss |
| from utils.image_utils import depth2normal |
| self.gaussians.update_learning_rate(step) |
| if step % 1000 == 0: |
| self.gaussians.oneupSHdegree() |
| scale = 1 |
| if step < 2000: |
| scale = 4 |
| elif step < 5000: |
| scale = 2 |
| if not self.viewpoint_stack: |
| self.viewpoint_stack = self.scene.getTrainCameras(scale).copy() |
| viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) |
| render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background, [float('inf'), float('inf')]) |
| image, normal, depth, opac = render_pkg["render"], render_pkg["normal"], render_pkg["depth"], render_pkg["opac"] |
| gt_image = viewpoint_cam.get_gtImage(self.background, self.dataset.use_mask) |
| mask_vis = (opac.detach() > 1e-5) |
| loss_target = (1.0 - self.opt.lambda_dssim) * l1_loss(image, gt_image) |
| loss_parasitic = self.opt.lambda_dssim * (1.0 - ssim(image, gt_image)) |
| d2n = depth2normal(depth, mask_vis, viewpoint_cam) |
| normal_norm = F.normalize(normal, dim=0) * mask_vis |
| loss_surface = cos_loss(normal_norm, d2n) |
| loss = loss_target + loss_parasitic |
| loss += (0.01 + 0.1 * min(2 * step / self.opt.iterations, 1)) * loss_surface |
| grad_cos_sim = 0.0 |
| p_ratio = 0.0 |
| |
| if self.track_decoupling and step % 100 == 0: |
| self.gaussians.optimizer.zero_grad() |
| loss_target.backward(retain_graph=True) |
| g_t = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else None |
| self.gaussians.optimizer.zero_grad() |
| loss_parasitic.backward(retain_graph=True) |
| g_p = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else None |
| if g_t is not None and g_p is not None: |
| mask = (g_t.norm(dim=1) > 1e-8) & (g_p.norm(dim=1) > 1e-8) |
| if mask.any(): |
| grad_cos_sim = float(F.cosine_similarity(g_t[mask], g_p[mask], dim=1).mean()) |
| p_ratio = float(g_p[mask].norm(dim=1).mean() / (g_t[mask].norm(dim=1).mean() + 1e-8)) |
| self.gaussians.optimizer.zero_grad() |
| loss.backward() |
| else: |
| loss.backward() |
|
|
| |
| with torch.no_grad(): |
| if step > self.opt.densify_from_iter: |
| self.gaussians.max_radii2D[render_pkg["visibility_filter"]] = torch.max(self.gaussians.max_radii2D[render_pkg["visibility_filter"]], render_pkg["radii"][render_pkg["visibility_filter"]]) |
| self.gaussians.add_densification_stats(render_pkg["viewspace_points"], render_pkg["visibility_filter"]) |
| if step % self.opt.densification_interval == 0: |
| self.gaussians.adaptive_prune(0.1, self.scene.cameras_extent) |
| if self.cap_gaussians is None or len(self.gaussians.get_xyz) < self.cap_gaussians: |
| self.gaussians.adaptive_densify(self.opt.densify_grad_threshold, self.scene.cameras_extent) |
| if step % self.opt.opacity_reset_interval == 0: |
| self.gaussians.reset_opacity(0.12, step) |
| |
| |
| self.gaussians.optimizer.step() |
| self.gaussians.optimizer.zero_grad(set_to_none=True) |
|
|
| num_gaussians = len(self.gaussians.get_xyz) |
| metrics = { |
| "loss": float(loss), |
| "loss_l1": float(loss_target), |
| "loss_ssim": float(loss_parasitic), |
| "num_gaussians": num_gaussians, |
| "delta_N": num_gaussians - self.last_n_gaussians, |
| "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024**3)), |
| "grad_cos_sim": grad_cos_sim, |
| "parasitic_ratio": p_ratio |
| } |
| self.last_n_gaussians = num_gaussians |
| |
| histograms = {} |
| if step % 1000 == 0: |
| histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).detach() |
| scales = torch.exp(self.gaussians._scaling).detach() |
| histograms["scaling"] = scales |
| scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales |
| 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): |
| from gaussian_renderer import render as native_render |
| with torch.no_grad(): |
| pkg = native_render(camera, self.gaussians, self.pipe, self.background, [float('inf'), float('inf')]) |
| return {"image": pkg["render"], "depth": pkg["depth"]} |
|
|
| 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, f"point_cloud/iteration_{iteration}/point_cloud.ply")) |
|
|
| def get_spatial_centers(self): |
| return self.gaussians._xyz.detach() |
|
|
| def compute_physical_metrics(self, cameras=None): |
| scales = self.gaussians.get_scaling.detach() |
| return { |
| "gamma_median": float(torch.median(scales[:, 0] / (scales[:, 1] + 1e-7))), |
| "alpha_mean": float(self.gaussians.get_opacity.detach().mean()), |
| "z_scale_mean": float(scales[:, 2].mean()) |
| } |
|
|
| def evaluate_spatial_field(self, query_points, cameras=None): |
| from pytorch3d.ops import knn_points |
| xyz = self.gaussians.get_xyz.detach() |
| opac = self.gaussians.get_opacity.detach() |
| dist_sq, _, _ = knn_points(query_points.unsqueeze(0), xyz.unsqueeze(0), K=1) |
| dist = torch.sqrt(dist_sq.squeeze(0)) |
| weights = torch.exp(-0.5 * (dist / 0.01)**2) |
| return (weights * opac.T).sum(dim=1) |
|
|
