| 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 |
|
|
| def force_unload_conflicts(): |
| conflicts = ['render', 'scene', 'utils', 'arguments', 'model'] |
| for name in list(sys.modules.keys()): |
| parts = name.split('.') |
| if parts[0] in conflicts: |
| del sys.modules[name] |
|
|
| @register_method("bgtriangle") |
| class BGTriangleWrapper(BaseMethod): |
| def __init__(self, dataset_config, hyperparams): |
| force_unload_conflicts() |
| |
| self.method_path = os.path.abspath("/root/autodl-tmp/bg_triangle_official") |
| if self.method_path not in sys.path: |
| sys.path.insert(0, self.method_path) |
| |
| try: |
| from utils.loss_utils import l1_loss, ssim |
| from arguments import ModelParams, PipelineParams, OptimizationParams |
| from model.bprimitive_bezier import BPrimitiveBezier |
| from render.renderer import Renderer |
| from scene import Scene |
| |
| self.l1_loss = l1_loss |
| self.ssim = ssim |
| |
| 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.order = hyperparams.get("order", 1) |
| self.bprimitive = BPrimitiveBezier(self.order, self.dataset.sh_degree) |
| self.bprimitive.boundary_mode = 1 |
| |
| |
| 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.bprimitive) |
| self.bprimitive.training_setup(self.opt) |
|
|
| N = self.bprimitive.control_point.shape[0] |
| if not hasattr(self.bprimitive, 'vis_map'): |
| self.bprimitive.vis_map = torch.zeros((N, 6, 6), device="cuda") |
| if not hasattr(self.bprimitive, 'vis_accum'): |
| self.bprimitive.vis_accum = torch.zeros((N, 1), device="cuda") |
| if not hasattr(self.bprimitive, 'denom_vis'): |
| self.bprimitive.denom_vis = torch.zeros((N, 1), device="cuda") |
| if not hasattr(self.bprimitive, 'edge_accum'): |
| self.bprimitive.edge_accum = torch.zeros((N, 1), device="cuda") |
| if not hasattr(self.bprimitive, 'denom_edge'): |
| self.bprimitive.denom_edge = torch.zeros((N, 1), device="cuda") |
| if not hasattr(self.bprimitive, 'gradient_accum'): |
| self.bprimitive.gradient_accum = torch.zeros((N, 1), device="cuda") |
| if not hasattr(self.bprimitive, 'denom'): |
| self.bprimitive.denom = torch.zeros((N, 1), device="cuda") |
|
|
| self.renderer = Renderer() |
| 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_primitives = self.bprimitive.control_point.shape[0] |
| |
| finally: |
| if self.method_path in sys.path: |
| sys.path.remove(self.method_path) |
|
|
| def train_iteration(self, step): |
| force_unload_conflicts() |
| if self.method_path not in sys.path: |
| sys.path.insert(0, self.method_path) |
| |
| 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)) |
|
|
| render_pkg, debug_info = self.renderer( |
| self.bprimitive, viewpoint_cam, self.background, None, |
| self.pipe.num_segments_per_bprimitive_edge, -5.0, 1.0, "version_1" |
| ) |
| image = render_pkg |
| |
| gt_image = viewpoint_cam.original_image.cuda() |
| loss_l1 = self.l1_loss(image, gt_image) |
| loss_ssim = 1.0 - self.ssim(image, gt_image) |
| loss_target = (1.0 - 0.2) * loss_l1 |
| loss_parasitic = 0.2 * loss_ssim |
| loss = loss_target + loss_parasitic |
|
|
| grad_cos_sim = 0.0 |
| parasitic_ratio = 0.0 |
|
|
| if self.track_decoupling and step % 100 == 0: |
| self.bprimitive.optimizer.zero_grad(set_to_none=True) |
| loss_target.backward(retain_graph=True) |
| |
| grad_target = self.bprimitive.control_point.grad.clone().reshape(self.bprimitive.control_point.shape[0], -1) if self.bprimitive.control_point.grad is not None else torch.zeros_like(self.bprimitive.control_point).reshape(self.bprimitive.control_point.shape[0], -1) |
| |
| self.bprimitive.optimizer.zero_grad(set_to_none=True) |
| loss_parasitic.backward(retain_graph=True) |
| |
| grad_parasitic = self.bprimitive.control_point.grad.clone().reshape(self.bprimitive.control_point.shape[0], -1) if self.bprimitive.control_point.grad is not None else torch.zeros_like(self.bprimitive.control_point).reshape(self.bprimitive.control_point.shape[0], -1) |
| |
| valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) |
| if valid_mask.any(): |
| grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) |
| parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) |
| |
| self.bprimitive.optimizer.zero_grad(set_to_none=True) |
| loss.backward() |
| else: |
| loss.backward() |
|
|
| with torch.no_grad(): |
| self.bprimitive.optimizer.step() |
| self.bprimitive.optimizer.zero_grad(set_to_none=True) |
|
|
| num_primitives = self.bprimitive.control_point.shape[0] |
| metrics = { |
| "loss": float(loss), |
| "loss_l1": float(loss_target), |
| "loss_ssim": float(loss_parasitic), |
| "num_gaussians": int(num_primitives), |
| "delta_N": int(num_primitives - self.last_n_primitives), |
| "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), |
| "grad_cos_sim": float(grad_cos_sim), |
| "parasitic_ratio": float(parasitic_ratio) |
| } |
| self.last_n_primitives = num_primitives |
|
|
| histograms = {} |
| if step % 1000 == 0: |
| histograms["opacity"] = torch.sigmoid(self.bprimitive.opacity).clone().detach() |
| scales = torch.exp(self.bprimitive.scaling).clone().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.bprimitive.control_point_dc.detach().norm(dim=-1).mean(dim=1) |
| |
| if self.method_path in sys.path: |
| sys.path.remove(self.method_path) |
| |
| return metrics, histograms |
|
|
| def render(self, camera): |
| force_unload_conflicts() |
| if self.method_path not in sys.path: |
| sys.path.insert(0, self.method_path) |
| |
| with torch.no_grad(): |
| image, _ = self.renderer( |
| self.bprimitive, camera, self.background, None, |
| self.pipe.num_segments_per_bprimitive_edge, -5.0, 1.0, "version_1" |
| ) |
| |
| if self.method_path in sys.path: |
| sys.path.remove(self.method_path) |
| |
| return {"image": image, "depth": None} |
|
|
| def save(self, save_dir, step): |
| force_unload_conflicts() |
| if self.method_path not in sys.path: |
| sys.path.insert(0, self.method_path) |
| |
| self.scene.save(step) |
| |
| if self.method_path in sys.path: |
| sys.path.remove(self.method_path) |
|
|
| def load(self, model_path, iteration): |
| pass |
|
|
| def get_spatial_centers(self): |
| return self.bprimitive.control_point.mean(dim=1) |
|
|
| def compute_physical_metrics(self, cameras=None): |
| metrics = {} |
| with torch.no_grad(): |
| raw_scales = self.bprimitive.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.bprimitive.opacity))) |
| |
| dc = self.bprimitive.control_point_dc |
| rest = self.bprimitive.control_point_rest |
| if rest is not None and rest.shape[2] > 0: |
| metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) |
|
|
| 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.bprimitive.control_point.mean(dim=1) |
| opacities = torch.sigmoid(self.bprimitive.opacity).squeeze() |
| scales = torch.exp(self.bprimitive.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 |
|
|
