| 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__), '../../CoAdaptGS'))) |
| from utils.loss_utils import l1_loss, ssim, SmoothLoss |
| from utils.graphics_utils import inverse_warp_images |
| from gaussian_renderer import render as native_render |
| from scene import Scene, GaussianModel |
| from arguments import ModelParams, PipelineParams, OptimizationParams |
|
|
| @register_method("coadaptgs") |
| class CoAdaptGSWrapper(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.parser.add_argument("--opacity_decay", action="store_true", default=True) |
| self.parser.add_argument("--opacity_decay_factor", type=float, default=0.995) |
| self.parser.add_argument("--dropout_factor", type=float, default=0.2) |
| self.parser.add_argument("--sigma_noise", type=float, default=0.8) |
| self.parser.add_argument("--cam_trans_dist", type=float, default=0.4) |
| self.parser.add_argument("--binocular_consistency", action="store_true", default=True) |
| self.parser.add_argument("--shift_cam_start", type=int, default=20000) |
| self.parser.add_argument("--n_views", type=int, default=-1) |
| self.parser.add_argument("--dataset_name", type=str, default="") |
| self.parser.add_argument("--suffix", type=str, default="") |
| |
| self.args, _ = self.parser.parse_known_args() |
| |
| self.args.source_path = dataset_config["source_path"] |
| self.args.model_path = dataset_config["model_path"] |
| self.args.eval = True |
| |
| self.args.n_views = 0 |
| self.args.dataset_name = "LLFF" |
| self.args.suffix = None |
| if not hasattr(self.args, 'images') or self.args.images is None: |
| self.args.images = "images" |
| 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) |
| |
| 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.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.smooth_loss = SmoothLoss() |
| |
| def train_iteration(self, step): |
| N = self.gaussians.get_xyz.shape[0] |
| if self.gaussians._scaling.shape[0] != N: |
| fallback_scale = float(self.gaussians._scaling.mean().item()) if self.gaussians._scaling.numel() > 0 else -4.0 |
| new_scales = torch.ones((N, 3), dtype=torch.float32, device="cuda") * fallback_scale |
| self.gaussians._scaling = torch.nn.Parameter(new_scales.requires_grad_(True)) |
| self.gaussians.training_setup(self.opt) |
|
|
| 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 = self.background |
| |
| render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, bg, |
| dropout_factor=self.args.dropout_factor, |
| sigma_noise=self.args.sigma_noise, train=True) |
| |
| image = render_pkg["render"] |
| viewspace_point_tensor = render_pkg["viewspace_points"] |
| visibility_filter = render_pkg["visibility_filter"] |
| radii = render_pkg["radii"] |
| depth = render_pkg.get("rendered_depth", render_pkg.get("depth", torch.zeros_like(image[0:1]))) |
| alpha = render_pkg.get("rendered_alpha", torch.zeros_like(image[0:1])) |
| dropout_mask = render_pkg.get("dropout_mask", torch.ones(self.gaussians.get_xyz.shape[0], dtype=torch.bool, device="cuda")) |
| |
| gt_image = viewpoint_cam.original_image.cuda() |
| |
| disparity_loss = torch.tensor(0.0, device="cuda") |
| if self.args.binocular_consistency and step > self.args.shift_cam_start: |
| trans_dist = (torch.rand(1) * self.args.cam_trans_dist * random.choice([-1.0, 1.0])).item() |
| shifted_cam = self.scene.getShiftedCamera(viewpoint_cam, trans_dist) |
| shifted_pkg = native_render(shifted_cam, self.gaussians, self.pipe, bg, dropout_factor=0.0, sigma_noise=0.0, train=True) |
| shifted_image = shifted_pkg["render"] |
| |
| focal_x = viewpoint_cam.get_focal()[0] if hasattr(viewpoint_cam, 'get_focal') else viewpoint_cam.focal_x |
| disparity = focal_x * (-trans_dist) / (depth + 1e-5) |
| |
| image_height = viewpoint_cam.image_height |
| image_width = viewpoint_cam.image_width |
| row_indices = torch.arange(0, image_height).view(-1, 1).repeat(1, image_width).cuda() |
| column_indices = torch.arange(0, image_width).repeat(image_height, 1).cuda() |
| mask_tensor = torch.ones((1, image_height, image_width), dtype=torch.float32).cuda() |
| |
| warped_image = inverse_warp_images(shifted_image.unsqueeze(0), disparity.unsqueeze(0), row_indices, column_indices) |
| shift_mask = inverse_warp_images(mask_tensor.unsqueeze(0), disparity.unsqueeze(0), row_indices, column_indices) |
| disparity_loss = l1_loss(warped_image, gt_image.unsqueeze(0), mask=shift_mask) + 0.05 * self.smooth_loss(disparity=disparity*shift_mask, image=gt_image.unsqueeze(0)) |
| |
| alpha_loss = torch.tensor(0.0, device="cuda") |
| if hasattr(viewpoint_cam, "gt_alpha_mask") and viewpoint_cam.gt_alpha_mask is not None: |
| alpha_loss = torch.mean(torch.abs(alpha) * (1 - viewpoint_cam.gt_alpha_mask)) |
| |
| Ll1 = l1_loss(image, gt_image) |
| loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 |
| loss_parasitic = self.opt.lambda_dssim * (1.0 - ssim(image, gt_image)) + disparity_loss + alpha_loss |
| loss = loss_target + loss_parasitic |
| |
| grad_cos_sim = 0.0 |
| parasitic_ratio = 0.0 |
| |
| if self.track_decoupling and step % 100 == 0: |
| self.gaussians.optimizer.zero_grad(set_to_none=True) |
| loss_target.backward(retain_graph=True) |
| grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) |
| |
| self.gaussians.optimizer.zero_grad(set_to_none=True) |
| loss_parasitic.backward(retain_graph=True) |
| grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) |
| |
| state = self.gaussians.optimizer.state.get(self.gaussians._xyz, None) |
| if state is not None and "exp_avg_sq" in state: |
| v_t = state["exp_avg_sq"] |
| eta_t = self.gaussians.optimizer.param_groups[0]["lr"] |
| u_target = (eta_t / (torch.sqrt(v_t) + 1e-8)) * grad_target |
| u_parasitic = (eta_t / (torch.sqrt(v_t) + 1e-8)) * grad_parasitic |
| else: |
| u_target = grad_target |
| u_parasitic = grad_parasitic |
| |
| valid_mask = (torch.norm(u_target, dim=1) > 0) & (torch.norm(u_parasitic, dim=1) > 0) |
| if valid_mask.any(): |
| grad_cos_sim = float(F.cosine_similarity(u_target[valid_mask], u_parasitic[valid_mask], dim=1).mean()) |
| parasitic_ratio = float(torch.norm(u_parasitic, dim=1).mean() / (torch.norm(u_target, dim=1).mean() + 1e-7)) |
| |
| self.gaussians.optimizer.zero_grad(set_to_none=True) |
| loss.backward() |
| else: |
| loss.backward() |
| |
| with torch.no_grad(): |
| if self.args.opacity_decay: |
| self.gaussians.opacity_decay(factor=self.args.opacity_decay_factor) |
| |
| if step < self.opt.densify_until_iter: |
| true_indices = torch.nonzero(dropout_mask, as_tuple=True)[0] |
| filtered_indices = true_indices[visibility_filter] |
| combined_mask = torch.zeros_like(dropout_mask, dtype=torch.bool).cuda() |
| combined_mask[filtered_indices] = True |
| self.gaussians.max_radii2D[combined_mask] = torch.max(self.gaussians.max_radii2D[combined_mask], radii[visibility_filter]) |
| self.gaussians.add_densification_stats(viewspace_point_tensor, combined_mask) |
| |
| if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: |
| size_threshold = None |
| self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold) |
| |
| self.gaussians.optimizer.step() |
| self.gaussians.optimizer.zero_grad(set_to_none=True) |
| |
| num_gaussians = self.gaussians.get_xyz.shape[0] |
| dropout_survival_ratio = float(dropout_mask.float().mean()) |
| visible_area_ratio = float(render_pkg.get("visi_area", torch.zeros(1)).float().mean()) |
| |
| metrics = { |
| "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic - disparity_loss - alpha_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), |
| "dropout_survival_ratio": dropout_survival_ratio, |
| "disparity_loss": float(disparity_loss), |
| "visible_area_ratio": visible_area_ratio |
| } |
| 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 |
| 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, dropout_factor=0.0, sigma_noise=0.0, train=False) |
| |
| image = render_pkg["render"] |
| depth = render_pkg.get("rendered_depth", render_pkg.get("depth", None)) |
| |
| if depth is None: |
| depth = torch.zeros_like(image[0:1]) |
| |
| if depth.dim() == 3: |
| depth = depth.unsqueeze(0) |
| |
| return {"image": image, "depth": 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, '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, opacities = self.gaussians._xyz, 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 |
|
|
