From 90957bcdfb288f1dc86dc45939e3895e534f091a Mon Sep 17 00:00:00 2001 From: Justin Kerr Date: Tue, 10 Oct 2023 17:22:42 -0700 Subject: [PATCH] cleanup --- nerfstudio/configs/method_configs.py | 4 +- nerfstudio/engine/trainer.py | 23 +- nerfstudio/models/gaussian_splatting.py | 1052 ++++++++--------- .../models/gaussian_splatting_inhouse.py | 704 ----------- nerfstudio/viewer_beta/control_panel.py | 5 +- nerfstudio/viewer_beta/render_panel.py | 1 + 6 files changed, 487 insertions(+), 1302 deletions(-) delete mode 100644 nerfstudio/models/gaussian_splatting_inhouse.py diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 1bf9324b5..50ff17e88 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -59,7 +59,7 @@ from nerfstudio.field_components.temporal_distortions import \ TemporalDistortionKind from nerfstudio.fields.sdf_field import SDFFieldConfig from nerfstudio.models.depth_nerfacto import DepthNerfactoModelConfig -from nerfstudio.models.gaussian_splatting_inhouse import GaussianSplattingModelConfig +from nerfstudio.models.gaussian_splatting import GaussianSplattingModelConfig from nerfstudio.models.generfacto import GenerfactoModelConfig from nerfstudio.models.instant_ngp import InstantNGPModelConfig from nerfstudio.data.dataparsers.colmap_dataparser import ColmapDataParserConfig @@ -643,7 +643,7 @@ method_configs["gaussian-splatting"] = TrainerConfig( }, }, viewer=ViewerConfig(num_rays_per_chunk=1 << 15), - vis="viewer", + vis="viewer_beta", ) diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index 3dede8071..6ce7862ed 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -267,21 +267,14 @@ class Trainer: # Skip the first two steps to avoid skewed timings that break the viewer rendering speed estimate. if step > 1: - if isinstance(self.pipeline.model, GaussianSplattingModel): - writer.put_scalar( - name=EventName.GAUSSIAN_NUM, - scalar=self.pipeline.model.get_gaussian_num, - step=step, - ) - else: - writer.put_time( - name=EventName.TRAIN_RAYS_PER_SEC, - duration=self.world_size - * self.pipeline.datamanager.get_train_rays_per_batch() - / max(0.001, train_t.duration), - step=step, - avg_over_steps=True, - ) + writer.put_time( + name=EventName.TRAIN_RAYS_PER_SEC, + duration=self.world_size + * self.pipeline.datamanager.get_train_rays_per_batch() + / max(0.001, train_t.duration), + step=step, + avg_over_steps=True, + ) self._update_viewer_state(step) diff --git a/nerfstudio/models/gaussian_splatting.py b/nerfstudio/models/gaussian_splatting.py index ba911fb54..4cf18c47a 100644 --- a/nerfstudio/models/gaussian_splatting.py +++ b/nerfstudio/models/gaussian_splatting.py @@ -20,318 +20,393 @@ from __future__ import annotations from dataclasses import dataclass, field from typing import Any, Dict, List, Literal, Optional, Tuple, Type, Union - -import matplotlib.pyplot as plt -import numpy as np import torch -import torch.nn as nn -from torch.cuda.amp.grad_scaler import GradScaler +from nerfstudio.data.scene_box import OrientedBox +from copy import deepcopy +from nerfstudio.cameras.rays import RayBundle from torch.nn import Parameter from torchmetrics.functional import structural_similarity_index_measure from torchmetrics.image import PeakSignalNoiseRatio from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity - -from nerfstudio.cameras.cameras import Cameras, CameraType -from nerfstudio.cameras.rays import RayBundle, RaySamples -from nerfstudio.engine.callbacks import (TrainingCallback, - TrainingCallbackAttributes, - TrainingCallbackLocation) +from diff_rast._torch_impl import compute_sh_color +from nerfstudio.cameras.cameras import Cameras +from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes, TrainingCallbackLocation +from nerfstudio.viewer_beta.viewer_elements import ViewerControl from nerfstudio.engine.optimizers import Optimizers -from nerfstudio.field_components.field_heads import FieldHeadNames -from nerfstudio.field_components.spatial_distortions import SceneContraction -from nerfstudio.fields.density_fields import HashMLPDensityField -from nerfstudio.fields.nerfacto_field import NerfactoField -from nerfstudio.model_components.losses import ( - MSELoss, distortion_loss, interlevel_loss, orientation_loss, - pred_normal_loss, scale_gradients_by_distance_squared) -from nerfstudio.model_components.ray_samplers import (ProposalNetworkSampler, - UniformSampler) -from nerfstudio.model_components.renderers import (AccumulationRenderer, - DepthRenderer, - NormalsRenderer, - RGBRenderer) -from nerfstudio.model_components.scene_colliders import NearFarCollider -from nerfstudio.model_components.shaders import NormalsShader from nerfstudio.models.base_model import Model, ModelConfig -from nerfstudio.utils import colormaps, gaussian_utils -from nerfstudio.utils import poses as pose_utils +import math +import numpy as np +from sklearn.neighbors import NearestNeighbors +import viser.transforms as vtf + + +from diff_rast.rasterize import RasterizeGaussians +from diff_rast.project_gaussians import ProjectGaussians +from diff_rast.sh import SphericalHarmonics, num_sh_bases + +def random_quat_tensor(N, **kwargs): + u = torch.rand(N, **kwargs) + v = torch.rand(N, **kwargs) + w = torch.rand(N, **kwargs) + return torch.stack( + [ + torch.sqrt(1 - u) * torch.sin(2 * math.pi * v), + torch.sqrt(1 - u) * torch.cos(2 * math.pi * v), + torch.sqrt(u) * torch.sin(2 * math.pi * w), + torch.sqrt(u) * torch.sin(2 * math.pi * w), + ], + dim=-1, + ) + +def RGB2SH(rgb): + C0 = 0.28209479177387814 + return (rgb - 0.5) / C0 + +def projection_matrix(znear, zfar, fovx, fovy, **kwargs): + t = znear * math.tan(0.5 * fovy) + b = -t + r = znear * math.tan(0.5 * fovx) + l = -r + n = znear + f = zfar + return torch.tensor( + [ + [2 * n / (r - l), 0.0, (r + l) / (r - l), 0.0], + [0.0, 2 * n / (t - b), (t + b) / (t - b), 0.0], + [0.0, 0.0, (f + n) / (f - n), -1.0 * f * n / (f - n)], + [0.0, 0.0, 1.0, 0.0], + ], + **kwargs, + ) @dataclass class GaussianSplattingModelConfig(ModelConfig): """Gaussian Splatting Model Config""" + _target: Type = field(default_factory=lambda: GaussianSplattingModel) - initial_ptc_size: int = 100000 - max_ptc_size: int = None - render_accumulation: bool = False - use_aabb: bool = True - initial_gaussian_scales: int = -3 - prune_density_after: int = 500 - prune_density_every: int = 100 - prune_density_until: int = 15000 - opacity_reset_interval: int = 3000 - one_up_sh_every: int = 1000 - lambda_ssim: float = 0.2 - use_diff_rast: bool = False # TODO (jake-austin): remove - gradient_threshold: float = 0.0002 - min_opacity: float = 0.005 - scene_extent: float = 2.0 - max_sh_degree: int = 3 + warmup_length:int = 1000 + """period of steps where refinement is turned off""" + refine_every: int = 100 + """period of steps where gaussians are culled and densified""" + resolution_schedule: int = 200 + """training starts at 1/d resolution, every n steps this is doubled""" + num_downscales: int = 2 + """at the beginning, resolution is 1/2^d, where d is this number""" + cull_alpha_thresh: float = 0.01 + """threshold of opacity for culling gaussians""" + cull_scale_thresh: float = 0.5 + """threshold of scale for culling gaussians""" + reset_alpha_every: int = 30 + """Every this many refinement steps, reset the alpha""" + densify_grad_thresh: float = 0.0002 + """threshold of positional gradient norm for densifying gaussians""" + densify_size_thresh: float = 0.01 + """below this size, gaussians are *duplicated*, otherwise split""" + sh_degree_interval: int = 1000 + """every n intervals turn on another sh degree""" + max_screen_size: int = 300 + """maximum screen size of a gaussian, in pixels""" + random_init: bool = False + """whether to initialize the positions uniformly randomly (not SFM points)""" + ssim_lambda: float = .2 + """weight of ssim loss""" + stop_split_at: int = 25000 + """stop splitting at this step""" + sh_degree: int = 4 + """maximum degree of spherical harmonics to use""" class GaussianSplattingModel(Model): """Gaussian Splatting model - TODO (jake-austin): Figure out how to print out on the training log in terminal the number of splats - Args: config: Gaussian Splatting configuration to instantiate model """ config: GaussianSplattingModelConfig + def __init__(self, *args, **kwargs): + if "seed_points" in kwargs: + self.seed_pts = kwargs["seed_points"] + else: + self.seed_pts = None + super().__init__(*args, **kwargs) + def populate_modules(self): - # TODO (jake-austin): clean this up, this is transplanted code across all the implementation functions - self.active_sh_degree = 0 - self.max_sh_degree = self.config.max_sh_degree + if self.seed_pts is not None and not self.config.random_init: + self.means = torch.nn.Parameter(self.seed_pts[0]) # (Location, Color) + else: + self.means = torch.nn.Parameter((torch.rand((100000, 3)) - 0.5) * 2) + self.xys_grad_norm = None + self.max_2Dsize = None + distances, _ = self.k_nearest_sklearn(self.means.data, 3) + distances = torch.from_numpy(distances) + # find the average of the three nearest neighbors for each point and use that as the scale + avg_dist = distances.mean(dim=-1, keepdim=True) + self.scales = torch.nn.Parameter(torch.log(avg_dist.repeat(1, 3))) - self._xyz = torch.empty(0) - self._features_dc = torch.empty(0) - self._features_rest = torch.empty(0) - self._scaling = torch.empty(0) - self._rotation = torch.empty(0) - self._opacity = torch.empty(0) - self.max_radii2D = torch.empty(0) - self.xyz_gradient_accum = torch.empty(0) - self.denom = torch.empty(0) - self.optimizer = None - self.percent_dense = 0 - self.spatial_lr_scale = 0 - self.max_points = self.config.max_ptc_size + self.quats = torch.nn.Parameter(random_quat_tensor(self.num_points)) + dim_sh = num_sh_bases(self.config.sh_degree) + if self.seed_pts is not None and not self.config.random_init: + fused_color = RGB2SH(self.seed_pts[1] / 255) + shs = torch.zeros((fused_color.shape[0], 3, dim_sh)).float().cuda() + shs[:, :3, 0] = fused_color + shs[:, 3:, 1:] = 0.0 + self.colors = torch.nn.Parameter(shs[:, :, 0:1]) + self.shs_rest = torch.nn.Parameter(shs[:, :, 1:]) + else: + self.colors = torch.nn.Parameter(torch.rand(self.num_points, 3, 1)) + self.shs_rest = torch.nn.Parameter(torch.zeros((self.num_points, 3, dim_sh - 1))) - def build_covariance_from_scaling_rotation(scaling, scaling_modifier, rotation): - L = gaussian_utils.build_scaling_rotation(scaling_modifier * scaling, rotation) - actual_covariance = L @ L.transpose(1, 2) - symm = gaussian_utils.strip_symmetric(actual_covariance) - return symm - - self.scaling_activation = torch.exp - self.scaling_inverse_activation = torch.log - - self.covariance_activation = build_covariance_from_scaling_rotation - - self.opacity_activation = torch.sigmoid - self.inverse_opacity_activation = gaussian_utils.inverse_sigmoid - - self.rotation_activation = torch.nn.functional.normalize - - num_pts = self.config.initial_ptc_size - print(f"Generating random point cloud ({num_pts})...") - - # We create random points inside the bounds of the synthetic Blender scenes - xyz = np.random.random((num_pts, 3)) * 2.6 - 1.3 - shs = np.random.random((num_pts, 3)) / 255.0 - # pcd = BasicPointCloud(points=xyz, colors=SH2RGB(shs), normals=np.zeros((num_pts, 3))) - - self.spatial_lr_scale = 5 - - fused_point_cloud = torch.tensor(np.asarray(xyz)).float().cuda() - fused_color = gaussian_utils.RGB2SH(torch.tensor(np.asarray(gaussian_utils.SH2RGB(shs))).float().cuda()) - features = torch.zeros((fused_color.shape[0], 3, (self.max_sh_degree + 1) ** 2)).float().cuda() - features[:, :3, 0 ] = fused_color - features[:, 3:, 1:] = 0.0 - - # dist2 = torch.clamp_min(distCUDA2(torch.from_numpy(np.asarray(pcd.points)).float().cuda()), 0.0000001) - scales = torch.ones((fused_point_cloud.shape[0], 3), device="cuda") * self.config.initial_gaussian_scales - rots = torch.zeros((fused_point_cloud.shape[0], 4), device="cuda") - rots[:, 0] = 1 - - opacities = gaussian_utils.inverse_sigmoid(0.1 * torch.ones((fused_point_cloud.shape[0], 1), dtype=torch.float, device="cuda")) - - self._xyz = nn.Parameter(fused_point_cloud.requires_grad_(True)) - self._features_dc = nn.Parameter(features[:,:,0:1].transpose(1, 2).contiguous().requires_grad_(True)) - self._features_rest = nn.Parameter(features[:,:,1:].transpose(1, 2).contiguous().requires_grad_(True)) - self._scaling = nn.Parameter(scales.requires_grad_(True)) - self._rotation = nn.Parameter(rots.requires_grad_(True)) - self._opacity = nn.Parameter(opacities.requires_grad_(True)) - self.max_radii2D = torch.zeros((self._xyz.shape[0]), device="cuda") - - self.percent_dense = 0.01 - self.xyz_gradient_accum = torch.zeros((self._xyz.shape[0], 1), device="cuda") - self.denom = torch.zeros((self._xyz.shape[0], 1), device="cuda") - - # l = [ - # {'params': [self._xyz], 'lr': training_args.position_lr_init*self.spatial_lr_scale, "name": "xyz"}, - # {'params': [self._features_dc], 'lr': training_args.feature_lr, "name": "f_dc"}, - # {'params': [self._features_rest], 'lr': training_args.feature_lr / 20.0, "name": "f_rest"}, - # {'params': [self._opacity], 'lr': training_args.opacity_lr, "name": "opacity"}, - # {'params': [self._scaling], 'lr': training_args.scaling_lr*self.spatial_lr_scale, "name": "scaling"}, - # {'params': [self._rotation], 'lr': training_args.rotation_lr, "name": "rotation"} - # ] - - # self.optimizer = torch.optim.Adam(l, lr=0.0, eps=1e-15) - # self.xyz_scheduler_args = gaussian_utils.get_expon_lr_func(lr_init=training_args.position_lr_init*self.spatial_lr_scale, - # lr_final=training_args.position_lr_final*self.spatial_lr_scale, - # lr_delay_mult=training_args.position_lr_delay_mult, - # max_steps=training_args.position_lr_max_steps) - - self.device_indicator_param = nn.Parameter(torch.empty(0)) - + self.opacities = torch.nn.Parameter(torch.logit(0.1 * torch.ones(self.num_points, 1))) # metrics self.psnr = PeakSignalNoiseRatio(data_range=1.0) - self.ssim = structural_similarity_index_measure + from torchmetrics.image import StructuralSimilarityIndexMeasure + + self.ssim = StructuralSimilarityIndexMeasure() self.lpips = LearnedPerceptualImagePatchSimilarity(normalize=True) + self.step = 0 + self.crop_box:Optional[OrientedBox] = None + self.back_color = torch.zeros(3) + + @property + def get_colors(self): + color = self.colors + shs_rest = self.shs_rest + return torch.cat((color, shs_rest), dim=2) + + def load_state_dict(self, dict, **kwargs): + # resize the parameters to match the new number of points + self.step = 30000 + newp = dict['means'].shape[0] + self.means = torch.nn.Parameter(torch.zeros(newp,3,device=self.device)) + self.scales = torch.nn.Parameter(torch.zeros(newp,3,device=self.device)) + self.quats = torch.nn.Parameter(torch.zeros(newp,4,device=self.device)) + self.colors = torch.nn.Parameter(torch.zeros(self.num_points, 3, 1,device=self.device)) + self.shs_rest = torch.nn.Parameter(torch.zeros(self.num_points, 3, num_sh_bases(self.config.sh_degree)-1,device=self.device)) + self.opacities = torch.nn.Parameter(torch.zeros(newp,1,device=self.device)) + super().load_state_dict(dict,**kwargs) + + def k_nearest_sklearn(self, x: torch.Tensor, k: int): + """ + Find k-nearest neighbors using sklearn's NearestNeighbors. + x: The data tensor of shape [num_samples, num_features] + k: The number of neighbors to retrieve + """ + # Convert tensor to numpy array + x_np = x.cpu().numpy() + + # Build the nearest neighbors model + nn_model = NearestNeighbors(n_neighbors=k + 1, algorithm="auto", metric="euclidean").fit(x_np) + + # Find the k-nearest neighbors + distances, indices = nn_model.kneighbors(x_np) + + # Exclude the point itself from the result and return + return distances[:, 1:].astype(np.float32), indices[:, 1:].astype(np.float32) + + def remove_from_optim(self, optimizer, deleted_mask, new_params): + """removes the deleted_mask from the optimizer provided""" + assert len(new_params) == 1 + assert isinstance(optimizer, torch.optim.Adam), "Only works with Adam" + + param = optimizer.param_groups[0]["params"][0] + param_state = optimizer.state[param] + del optimizer.state[param] + + # Modify the state directly without deleting and reassigning. + param_state['exp_avg'] = param_state['exp_avg'][~deleted_mask] + param_state['exp_avg_sq'] = param_state['exp_avg_sq'][~deleted_mask] + + # Update the parameter in the optimizer's param group. + del optimizer.param_groups[0]["params"][0] + del optimizer.param_groups[0]["params"] + optimizer.param_groups[0]["params"] = new_params + optimizer.state[new_params[0]] = param_state + + def dup_in_optim(self,optimizer, dup_mask, new_params,n=2): + """adds the parameters to the optimizer""" + param = optimizer.param_groups[0]["params"][0] + param_state = optimizer.state[param] + repeat_dims = (n,) + tuple(1 for _ in range(param_state['exp_avg'].dim() - 1)) + param_state['exp_avg'] = torch.cat([param_state['exp_avg'], torch.zeros_like(param_state['exp_avg'][dup_mask.squeeze()]).repeat(*repeat_dims)],dim=0) + param_state['exp_avg_sq'] = torch.cat([param_state['exp_avg_sq'], torch.zeros_like(param_state['exp_avg_sq'][dup_mask.squeeze()]).repeat(*repeat_dims)],dim=0) + del optimizer.state[param] + optimizer.state[new_params[0]] = param_state + optimizer.param_groups[0]["params"] = new_params + del param + + def after_train(self, step): + with torch.no_grad(): + # keep track of a moving average of grad norms + visible_mask = (self.radii > 0).flatten() + grads = self.xys.grad.detach().norm(dim=-1)#TODO fill in + # print(f"grad norm min {grads.min().item()} max {grads.max().item()} mean {grads.mean().item()} size {grads.shape}") + if self.xys_grad_norm is None: + self.xys_grad_norm = grads + self.vis_counts = torch.ones_like(self.xys_grad_norm) + else: + self.vis_counts[visible_mask] = self.vis_counts[visible_mask] + 1 + self.xys_grad_norm[visible_mask] = grads[visible_mask] + self.xys_grad_norm[visible_mask] + + # update the max screen size + if self.max_2Dsize is None: + self.max_2Dsize = torch.zeros_like(self.radii) + newradii = self.radii[visible_mask] + self.max_2Dsize[visible_mask] = torch.maximum(self.max_2Dsize[visible_mask], newradii) + + def set_crop(self,crop_box:OrientedBox): + self.crop_box = crop_box + + def refinement_after(self, optimizers: Optimizers, step): + if self.step > self.config.warmup_length: + with torch.no_grad(): + if self.step self.config.densify_grad_thresh).squeeze() + splits = (self.scales.exp().max(dim=-1).values > self.config.densify_size_thresh).squeeze() + splits &= high_grads + nsamps=2 + split_means,split_colors,split_shs,split_opacities,split_scales,split_quats = self.split_gaussians(splits,nsamps) + + dups = (self.scales.exp().max(dim=-1).values <= self.config.densify_size_thresh).squeeze() + dups &= high_grads + dup_means,dup_colors,dup_shs,dup_opacities,dup_scales,dup_quats = self.dup_gaussians(dups) + self.means = Parameter(torch.cat([self.means.detach(),split_means, dup_means],dim=0)) + self.colors = Parameter(torch.cat([self.colors.detach(),split_colors, dup_colors],dim=0)) + self.shs_rest = Parameter(torch.cat([self.shs_rest.detach(),split_shs, dup_shs],dim=0)) + self.opacities = Parameter(torch.cat([self.opacities.detach(),split_opacities, dup_opacities],dim=0)) + self.scales = Parameter(torch.cat([self.scales.detach(),split_scales, dup_scales],dim=0)) + self.quats = Parameter(torch.cat([self.quats.detach(),split_quats, dup_quats],dim=0)) + #append zeros to the max_2Dsize tensor + self.max_2Dsize = torch.cat([self.max_2Dsize,torch.zeros_like(split_scales[:,0]),torch.zeros_like(dup_scales[:,0])],dim=0) + split_idcs = torch.where(splits)[0] + param_groups = self.get_param_groups() + for group,param in param_groups.items(): + self.dup_in_optim(optimizers.optimizers[group],split_idcs,param,n=nsamps) + dup_idcs = torch.where(dups)[0] + + param_groups = self.get_param_groups() + for group,param in param_groups.items(): + self.dup_in_optim(optimizers.optimizers[group],dup_idcs,param,1) + + #only cull if we've seen every image since opacity reset + reset_interval = (self.config.reset_alpha_every * self.config.refine_every) + if self.step % reset_interval > self.num_train_data: + #then cull + deleted_mask = self.cull_gaussians() + param_groups = self.get_param_groups() + for group, param in param_groups.items(): + self.remove_from_optim(optimizers.optimizers[group], deleted_mask, param) + + if self.step // self.config.refine_every % self.config.reset_alpha_every == 0: + reset_value = 0.01 + self.opacities.data = torch.full_like( + self.opacities.data, torch.logit(torch.tensor(reset_value)).item() + ) + # reset the exp of optimizer + optim = optimizers.optimizers["opacity"] + param = optim.param_groups[0]["params"][0] + param_state = optim.state[param] + param_state["exp_avg"] = torch.zeros_like(param_state["exp_avg"]) + param_state["exp_avg_sq"] = torch.zeros_like(param_state["exp_avg_sq"]) + self.xys_grad_norm = None + self.vis_counts = None + self.max_2Dsize = None + + def cull_gaussians(self): + """ + This function deletes gaussians with under a certain opacity threshold + """ + n_bef = self.num_points + # cull transparent ones + culls = (torch.sigmoid(self.opacities) < self.config.cull_alpha_thresh).squeeze() + if self.step > self.config.refine_every * self.config.reset_alpha_every: + # cull huge ones that are in the center of the scene (a 2-unit box is twice as large as all the cameras) + toobigs = (torch.exp(self.scales).max(dim=-1).values > self.config.cull_scale_thresh).squeeze() & (self.means.amax(dim=-1) < 2) + culls = culls | toobigs + # cull big screen space + culls = culls | (self.max_2Dsize > self.config.max_screen_size).squeeze() + self.means = Parameter(self.means[~culls].detach()) + self.scales = Parameter(self.scales[~culls].detach()) + self.quats = Parameter(self.quats[~culls].detach()) + self.colors = Parameter(self.colors[~culls].detach()) + self.shs_rest = Parameter(self.shs_rest[~culls].detach()) + self.opacities = Parameter(self.opacities[~culls].detach()) + + print(f"Culled {n_bef - self.num_points} gaussians") + return culls + + def split_gaussians(self, split_mask,samps): + """ + This function splits gaussians that are too large + """ + from diff_rast._torch_impl import quat_to_rotmat + n_splits = split_mask.sum().item() + print(f"Splitting {split_mask.sum().item()/self.num_points} gaussians: {n_splits}/{self.num_points}") + centered_samples = torch.randn((samps*n_splits,3),device=self.device) # Nx3 of axis-aligned scales + scaled_samples = torch.exp(self.scales[split_mask].repeat(samps,1)) * centered_samples # how these scales are rotated + quats = self.quats[split_mask]/self.quats[split_mask].norm(dim=-1,keepdim=True) #normalize them first + rots = quat_to_rotmat(quats.repeat(samps,1)) # how these scales are rotated + rotated_samples = torch.bmm(rots,scaled_samples[...,None]).squeeze() + new_means = rotated_samples + self.means[split_mask].repeat(samps,1) + # step 2, sample new colors + new_colors = self.colors[split_mask].repeat(samps,1,1) + # step 3, sample new shs + new_shs = self.shs_rest[split_mask].repeat(samps,1,1) + # step 3, sample new opacities + new_opacities = self.opacities[split_mask].repeat(samps,1) + # step 4, sample new scales + new_scales = torch.log(torch.exp(self.scales[split_mask])/1.6).repeat(samps,1) + self.scales[split_mask] = torch.log(torch.exp(self.scales[split_mask])/1.6) + # step 5, sample new quats + new_quats = self.quats[split_mask].repeat(samps,1) + return new_means,new_colors,new_shs,new_opacities,new_scales,new_quats + + def dup_gaussians(self, dup_mask): + """ + This function duplicates gaussians that are too small + """ + n_dups = dup_mask.sum().item() + print(f"Would duplicate {dup_mask.sum().item()/self.num_points} gaussians: {n_dups}/{self.num_points}") + dup_means = self.means[dup_mask] + dup_colors = self.colors[dup_mask] + dup_shs = self.shs_rest[dup_mask] + dup_opacities = self.opacities[dup_mask] + dup_scales = self.scales[dup_mask] + dup_quats = self.quats[dup_mask] + return dup_means, dup_colors, dup_shs, dup_opacities, dup_scales, dup_quats @property - def get_scaling(self): - return self.scaling_activation(self._scaling) - - @property - def get_rotation(self): - return self.rotation_activation(self._rotation) - - @property - def get_gaussian_num(self): - return self._xyz.shape[0] - - @property - def get_xyz(self): - return self._xyz - - @property - def get_features(self): - features_dc = self._features_dc - features_rest = self._features_rest - - return torch.cat((features_dc, features_rest), dim=1) - - @property - def get_opacity(self): - return self.opacity_activation(self._opacity) - - def get_covariance(self, scaling_modifier = 1): - return self.covariance_activation(self.get_scaling, scaling_modifier, self._rotation) - - @property - def device(self): - """Returns the device that the model is on.""" - return self.device_indicator_param.device + def num_points(self): + return self.means.shape[0] def get_training_callbacks( self, training_callback_attributes: TrainingCallbackAttributes ) -> List[TrainingCallback]: - - # TODO (jake-austin): remove and implement this - callbacks = [] - - def oneupSHdegree(step): - if step > 0 and self.active_sh_degree < self.max_sh_degree: - self.active_sh_degree += 1 - - def wrapper_add_densification_stats(step): - if step < self.config.prune_density_until: - self.add_densification_stats(self.densification_stats_cache[0], self.densification_stats_cache[1]) - - callbacks.append( + cbs = [] + cbs.append(TrainingCallback([TrainingCallbackLocation.BEFORE_TRAIN_ITERATION], self.step_cb)) + # The order of these matters + cbs.append( TrainingCallback( - where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION], - update_every_num_iters=1, - func=wrapper_add_densification_stats, + [TrainingCallbackLocation.AFTER_TRAIN_ITERATION], + self.after_train, ) ) - - def wrapper_pruning_densification(step): - if step >= self.config.prune_density_after and step < self.config.prune_density_until: - size_threshold = None#20 if iteration > opt.opacity_reset_interval else None - self.optimizer = training_callback_attributes.optimizers - self.densify_and_prune(self.config.gradient_threshold, self.config.min_opacity, self.config.scene_extent, None) - - def wrapper_reset_opacity(step): - if step < self.config.prune_density_until and step > 0: - self.reset_opacity() - - callbacks.append( + cbs.append( TrainingCallback( - where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION], - update_every_num_iters=self.config.one_up_sh_every, - func=oneupSHdegree, - ) - ) - callbacks.append( - TrainingCallback( - where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION], - update_every_num_iters=self.config.prune_density_every, - func=wrapper_pruning_densification, - ) - ) - callbacks.append( - TrainingCallback( - where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION], - update_every_num_iters=self.config.opacity_reset_interval, - func=wrapper_reset_opacity, + [TrainingCallbackLocation.AFTER_TRAIN_ITERATION], + self.refinement_after, + update_every_num_iters=self.config.refine_every, + args=[training_callback_attributes.optimizers], ) ) + return cbs - return callbacks - - # def pruning_densification(self, optimizers: Optimizers, grad_scaler: GradScaler): - - # prune_mask = (self.get_opacity < 0.005).squeeze() - # # if max_screen_size: - # # big_points_vs = self.max_radii2D > max_screen_size - # # big_points_ws = self.get_scaling.max(dim=1).values > 0.1 * extent - # # prune_mask = torch.logical_or(torch.logical_or(prune_mask, big_points_vs), big_points_ws) - # self.prune_points(prune_mask) - - # mask = torch.cat((selected_pts_mask, torch.zeros(N * selected_pts_mask.sum(), device="cuda", dtype=bool))) - - # valid_points_mask = ~mask - # for group_name in self.optimizers.optimizers: - # stored_state = self.optimizer.state.get(group['params'][0], None) - # if stored_state is not None: - # stored_state["exp_avg"] = stored_state["exp_avg"][mask] - # stored_state["exp_avg_sq"] = stored_state["exp_avg_sq"][mask] - - # del self.optimizer.state[group['params'][0]] - # group["params"][0] = nn.Parameter((group["params"][0][mask].requires_grad_(True))) - # self.optimizer.state[group['params'][0]] = stored_state - - # optimizable_tensors[group["name"]] = group["params"][0] - # else: - # group["params"][0] = nn.Parameter(group["params"][0][mask].requires_grad_(True)) - # optimizable_tensors[group["name"]] = group["params"][0] - - # self._xyz = optimizable_tensors["xyz"] - # self._features_dc = optimizable_tensors["f_dc"] - # self._features_rest = optimizable_tensors["f_rest"] - # self._opacity = optimizable_tensors["opacity"] - # self._scaling = optimizable_tensors["scaling"] - # self._rotation = optimizable_tensors["rotation"] - - # self.xyz_gradient_accum = self.xyz_gradient_accum[valid_points_mask] - - # self.denom = self.denom[valid_points_mask] - # self.max_radii2D = self.max_radii2D[valid_points_mask] - - - # def reset_opacity(self, optimizers, grad_scaler): - # opacities_new = gaussian_utils.inverse_sigmoid(torch.min(self.get_opacity, torch.ones_like(self.get_opacity)*0.01)) - # optimizable_tensors = self.update_optimizer_tensors({"opacity": opacities_new}, "replace") - # self._opacity = optimizable_tensors["opacity"] - - # def update_optimizer_tensors(self, data: dict, operation: Literal["cat", "prune", "replace"]): - # optimizable_tensors = {} - # for group in self.optimizer.param_groups: - # if group["name"] in data: - # if operation == "cat": - # # TODO (jake-austin): implement this, consider splitting this into 3 functions? - - # elif operation == "prune": - - # elif operation == "replace": - # assert False - # else: - # raise ValueError(f"Unknown operation {operation}") + def step_cb(self, step): + self.step = step def get_param_groups(self) -> Dict[str, List[Parameter]]: """Obtain the parameter groups for the optimizers @@ -340,18 +415,19 @@ class GaussianSplattingModel(Model): Mapping of different parameter groups """ return { - "xyz": [self._xyz], - "f_dc": [self._features_dc], - "f_rest": [self._features_rest], - "opacity": [self._opacity], - "scaling": [self._scaling], - "rotation": [self._rotation], + "xyz": [self.means], + "color": [self.colors], + "shs": [self.shs_rest], + "opacity": [self.opacities], + "scaling": [self.scales], + "rotation": [self.quats], } - def get_outputs(self, ray_bundle: Cameras) -> Dict[str, Union[torch.Tensor, List]]: - """Takes in a Ray Bundle and returns a dictionary of outputs. + def _get_downscale_factor(self): + return 2 ** max((self.config.num_downscales - self.step // self.config.resolution_schedule), 0) - TODO (jake-austin): use the new homebrew nerfstudio gaussian rasterization code instead + def get_outputs(self, camera: Cameras) -> Dict[str, Union[torch.Tensor, List]]: + """Takes in a Ray Bundle and returns a dictionary of outputs. Args: ray_bundle: Input bundle of rays. This raybundle should have all the @@ -360,122 +436,111 @@ class GaussianSplattingModel(Model): Returns: Outputs of model. (ie. rendered colors) """ - - #print(f"Num Points: {self._xyz.shape[0]}") - - - camera = ray_bundle.reshape(()) - assert camera.shape == () - R = camera.camera_to_worlds[..., :3, :3].squeeze() - R[:,0] = -R[:,0] - T = (R.T @ camera.camera_to_worlds[..., :3, 3:4]).squeeze() - fovx = gaussian_utils.focal2fov(camera.fx, camera.width) - fovy = gaussian_utils.focal2fov(camera.fy, camera.height) - - world_view_transform = gaussian_utils.getWorld2View2(R, T, torch.tensor([0.0, 0.0, 0.0]).to(self.device), 1.0).transpose(0, 1).to(self.device) - projection_matrix = gaussian_utils.getProjectionMatrix(znear=.01, zfar=100, fovX=fovx, fovY=fovy).transpose(0,1).to(self.device) - full_projection_matrix = (world_view_transform.unsqueeze(0).bmm(projection_matrix.unsqueeze(0))).squeeze(0) - - if not self.config.use_diff_rast: - - camera_center = world_view_transform.inverse()[3, :3] - data = { - "xyz": self.get_xyz, - "features": self.get_features, - "opacity": self.get_opacity, - "scaling": self.get_scaling, - "rotation": self.get_rotation, - "active_sh_degree": self.active_sh_degree, - "max_sh_degree": self.max_sh_degree, - "FoVx": fovx, - "FoVy": fovy, - "image_height": camera.height, - "image_width": camera.width, - "world_view_transform": world_view_transform, - "full_proj_transform": full_projection_matrix, - "camera_center": camera_center, - } - - outs = gaussian_utils.render_from_dict( - data, - bg_color=torch.tensor([1,1,1]).to(torch.float32).to(self.device), # TODO (jake-austin): This should be in config - ) #HWC - image = outs["render"].permute(1, 2, 0) - - if self.train(): - self.densification_stats_cache = [outs["viewspace_points"], outs["visibility_filter"]] - - return {"rgb": image} - + if not isinstance(camera, Cameras): + print("Called get_outputs with not a camera") + return {} + assert camera.shape[0] == 1, "Only one camera at a time" + # dont mutate the input + camera = deepcopy(camera) + if self.training: + d = self._get_downscale_factor() + camera.rescale_output_resolution(1/d) + if self.crop_box is not None and not self.training: + crop_ids = self.crop_box.within(self.means).squeeze() + if crop_ids.sum()==0: + return {'rgb':torch.full((camera.height.item(),camera.width.item(),3),.5,device=self.device)} else: - - c2w = camera.camera_to_worlds - c2w = pose_utils.to4x4(c2w) - w2c = torch.linalg.inv(c2w) - proj = self._get_proj_matrix( - w2c=w2c, - fx=camera.fx.item(), - fy=camera.fy.item(), - width=camera.width.item(), - height=camera.height.item(), - ) - - BLOCK_X = 16 - BLOCK_Y = 16 - W = camera.width.item() - H = camera.height.item() - tile_bounds = (W + BLOCK_X - 1) // BLOCK_X, (H + BLOCK_Y - 1) // BLOCK_Y, 1 - - glob_scale = 1 - xys, depths, radii, conics, num_tiles_hit = diff_rast.project_gaussians( - means3d=self.get_xyz, - scales=self.get_scaling, - glob_scale=glob_scale, - quats=self.get_rotation, - viewmat=w2c, - projmat=proj, - img_height=camera.width.item(), - img_width=camera.height.item(), - fx=camera.fx.item(), - fy=camera.fy.item(), - tile_bounds=tile_bounds, - ) - - rgb = diff_rast.rasterize( - xys, - depths, - radii, - conics, - num_tiles_hit, - self.get_features, - self.get_opacity.squeeze(), - camera.height.item(), - camera.width.item(), - ) - - return {"rgb": rgb} - - def _get_proj_matrix( - self, w2c, fx: float, fy: float, width: int, height: int, znear=0.01, zfar=100 - ): - top = 0.5 * height / fy * znear - bottom = -top - right = 0.5 * width / fx * znear - left = -right - - P = torch.zeros(4, 4).to(self.device) - - z_sign = 1.0 - - P[0, 0] = 2.0 * znear / (right - left) - P[1, 1] = 2.0 * znear / (top - bottom) - P[0, 2] = (right + left) / (right - left) - P[1, 2] = (top + bottom) / (top - bottom) - P[3, 2] = z_sign - P[2, 2] = z_sign * zfar / (zfar - znear) - P[2, 3] = -(zfar * znear) / (zfar - znear) - - return P @ w2c + crop_ids = torch.ones_like(self.means[:,0],dtype=torch.bool) + #shift the camera to center of scene looking at center + R = camera.camera_to_worlds[..., :3, :3] # 1 x 3 x 3 + T = camera.camera_to_worlds[..., :3, 3:4] # 1 x 3 x 1 + R = vtf.SO3.from_matrix(R.cpu().squeeze().numpy()) + R = R @ vtf.SO3.from_x_radians(np.pi) + R = torch.from_numpy(R.as_matrix()[None, ...]).to(self.device, torch.float32) + viewmat = torch.cat([R, T], dim=2) + # add a row of zeros and a 1 to the bottom of the viewmat + viewmat = torch.cat([viewmat, torch.tensor([[[0, 0, 0, 1]]], device=self.device)], dim=1) + # invert it + viewmat = torch.inverse(viewmat) + # calculate the FOV of the camera given fx and fy, width and height + cx = camera.cx.item() + cy = camera.cy.item() + fovx = 2 * math.atan(camera.width / (2 * camera.fx)) + fovy = 2 * math.atan(camera.height / (2 * camera.fy)) + W, H = camera.width.item(), camera.height.item() + self.last_size = (H, W) + projmat = projection_matrix(0.001, 1000, fovx, fovy).to(self.device) + BLOCK_X, BLOCK_Y = 16, 16 + tile_bounds = ( + (W + BLOCK_X - 1) // BLOCK_X, + (H + BLOCK_Y - 1) // BLOCK_Y, + 1, + ) + if self.training: + background = torch.rand(3, device=self.device) + else: + background = torch.zeros(3, device=self.device) + + opacities_crop = self.opacities[crop_ids] + means_crop = self.means[crop_ids] + colors_crop = self.get_colors[crop_ids] + self.xys, depths, self.radii, conics, num_tiles_hit, cov3d = ProjectGaussians.apply( + self.means[crop_ids], + torch.exp(self.scales[crop_ids]), + 1, + self.quats[crop_ids]/self.quats[crop_ids].norm(dim=-1,keepdim=True), + viewmat.squeeze()[:3, :], + projmat.squeeze() @ viewmat.squeeze(), + camera.fx.item(), + camera.fy.item(), + cx, + cy, + H, + W, + tile_bounds, + ) + #Important to allow xys grads to populate properly + if self.training: + self.xys.retain_grad() + rend_mask = self.radii>0 + if self.config.sh_degree > 0: + viewdirs = means_crop[rend_mask].detach() - camera.camera_to_worlds[..., :3, 3] # (N, 3) + viewdirs = viewdirs / viewdirs.norm(dim=-1, keepdim=True) + n = min(self.step // self.config.sh_degree_interval,self.config.sh_degree) + # rgbs = eval_sh(n, self.get_colors, viewdirs)[crop_ids] + n_coeffs = num_sh_bases(n) + coeffs = colors_crop.permute(0,2,1)[rend_mask,:n_coeffs,:] + rgbs = SphericalHarmonics.apply(n, viewdirs,coeffs) + rgbs = torch.clamp(rgbs + 0.5,0.0,1.0) + else: + rgbs = self.get_colors.squeeze()[rend_mask] # (N, 3) + rgbs = torch.sigmoid(rgbs) + + rgb = RasterizeGaussians.apply( + self.xys[rend_mask], + depths[rend_mask], + self.radii[rend_mask], + conics[rend_mask], + num_tiles_hit[rend_mask], + rgbs, + torch.sigmoid(opacities_crop[rend_mask]), + H, + W, + background, + ) + depth_im = RasterizeGaussians.apply( + self.xys[rend_mask], + depths[rend_mask], + self.radii[rend_mask], + conics[rend_mask], + num_tiles_hit[rend_mask], + depths[:,None], + torch.sigmoid(opacities_crop[rend_mask]), + H, + W, + torch.ones(1,device=self.device)*10, + )[...,0:1] + return {"rgb": rgb[..., :3],'depth':depth_im} def get_metrics_dict(self, outputs, batch) -> Dict[str, torch.Tensor]: """Compute and returns metrics. @@ -485,7 +550,7 @@ class GaussianSplattingModel(Model): batch: ground truth batch corresponding to outputs """ - return {"Num Gaussians": self._xyz.shape[0]} + return {} def get_loss_dict(self, outputs, batch, metrics_dict=None) -> Dict[str, torch.Tensor]: """Computes and returns the losses dict. @@ -495,21 +560,30 @@ class GaussianSplattingModel(Model): batch: ground truth batch corresponding to outputs metrics_dict: dictionary of metrics, some of which we can use for loss """ - # plt.imsave("test.png", np.concatenate([outputs["rgb"].detach().cpu().numpy(), batch["image"].detach().cpu().numpy()], axis=1)) + d = self._get_downscale_factor() + if d > 1: + # use torchvision to resize + import torchvision.transforms.functional as TF - Ll1 = gaussian_utils.l1_loss(batch['image'], outputs['rgb']) - loss = (1.0 - self.config.lambda_ssim) * Ll1 + self.config.lambda_ssim * (1.0 - gaussian_utils.ssim(batch['image'], outputs['rgb'])) - return {"main_loss": loss} + newsize = (batch["image"].shape[0] // d, batch["image"].shape[1] // d) + gt_img = TF.resize(batch["image"].permute(2, 0, 1), newsize).permute(1, 2, 0) + else: + gt_img = batch['image'] + Ll1 = torch.abs(gt_img- outputs['rgb']).mean() + simloss = (1-self.ssim(gt_img.permute(2,0,1)[None,...], outputs['rgb'].permute(2,0,1)[None,...])) + return {"main_loss": (1-self.config.ssim_lambda)*Ll1 + self.config.ssim_lambda*simloss} @torch.no_grad() - def get_outputs_for_camera_ray_bundle(self, camera_ray_bundle: Cameras) -> Dict[str, torch.Tensor]: + def get_outputs_for_camera_ray_bundle( + self, camera_ray_bundle: RayBundle, camera: Optional[Cameras] = None + ) -> Dict[str, torch.Tensor]: """Takes in camera parameters and computes the output of the model. Args: camera_ray_bundle: ray bundle to calculate outputs over """ - outs = self.get_outputs(camera_ray_bundle.to(self.device)) - outs["rgb"] = torch.clamp(outs["rgb"], 0.0, 1.0) + assert camera is not None, "must provide camera to gaussian model" + outs = self.get_outputs(camera.to(self.device)) return outs def get_image_metrics_and_images( @@ -546,179 +620,3 @@ class GaussianSplattingModel(Model): images_dict = {"img": combined_rgb} return metrics_dict, images_dict - - - def replace_tensor_to_optimizer(self, tensor, name): - optimizable_tensors = {} - for group_k, group_v in self.optimizer.optimizers.items(): - if group_k == name: - stored_state = self.optimizer.optimizers[group_k].state.get(group_v.param_groups[0]['params'][0], None) - stored_state["exp_avg"] = torch.zeros_like(tensor) - stored_state["exp_avg_sq"] = torch.zeros_like(tensor) - - del self.optimizer.optimizers[group_k].state[group_v.param_groups[0]['params'][0]] - group_v.param_groups[0]["params"][0] = nn.Parameter(tensor.requires_grad_(True)) - self.optimizer.optimizers[group_k].state[group_v.param_groups[0]['params'][0]] = stored_state - - optimizable_tensors[group_k] = group_v.param_groups[0]["params"][0] - return optimizable_tensors - - def _prune_optimizer(self, mask): - optimizable_tensors = {} - for group_k, group_v in self.optimizer.optimizers.items(): - stored_state = self.optimizer.optimizers[group_k].state.get(group_v.param_groups[0]['params'][0], None) - if stored_state is not None: - stored_state["exp_avg"] = stored_state["exp_avg"][mask] - stored_state["exp_avg_sq"] = stored_state["exp_avg_sq"][mask] - - del self.optimizer.optimizers[group_k].state[group_v.param_groups[0]['params'][0]] - group_v.param_groups[0]["params"][0] = nn.Parameter((group_v.param_groups[0]["params"][0][mask].requires_grad_(True))) - self.optimizer.optimizers[group_k].state[group_v.param_groups[0]['params'][0]] = stored_state - - optimizable_tensors[group_k] = group_v.param_groups[0]["params"][0] - else: - group_v.param_groups[0]["params"][0] = nn.Parameter(group_v.param_groups[0]["params"][0][mask].requires_grad_(True)) - optimizable_tensors[group_k] = group_v.param_groups[0]["params"][0] - return optimizable_tensors - - def prune_points(self, mask): - valid_points_mask = ~mask - optimizable_tensors = self._prune_optimizer(valid_points_mask) - - self._xyz = optimizable_tensors["xyz"] - self._features_dc = optimizable_tensors["f_dc"] - self._features_rest = optimizable_tensors["f_rest"] - self._opacity = optimizable_tensors["opacity"] - self._scaling = optimizable_tensors["scaling"] - self._rotation = optimizable_tensors["rotation"] - - self.xyz_gradient_accum = self.xyz_gradient_accum[valid_points_mask] - - self.denom = self.denom[valid_points_mask] - self.max_radii2D = self.max_radii2D[valid_points_mask] - - def cat_tensors_to_optimizer(self, tensors_dict): - optimizable_tensors = {} - for group_k, group_v in self.optimizer.optimizers.items(): - assert len(group_v.param_groups[0]["params"]) == 1 - extension_tensor = tensors_dict[group_k] - stored_state = self.optimizer.optimizers[group_k].state.get(group_v.param_groups[0]['params'][0], None) - if stored_state is not None: - stored_state["exp_avg"] = torch.cat((stored_state["exp_avg"], torch.zeros_like(extension_tensor)), dim=0) - stored_state["exp_avg_sq"] = torch.cat((stored_state["exp_avg_sq"], torch.zeros_like(extension_tensor)), dim=0) - - del self.optimizer.optimizers[group_k].state[group_v.param_groups[0]['params'][0]] - group_v.param_groups[0]["params"][0] = nn.Parameter(torch.cat((group_v.param_groups[0]["params"][0], extension_tensor), dim=0).requires_grad_(True)) - self.optimizer.optimizers[group_k].state[group_v.param_groups[0]['params'][0]] = stored_state - - optimizable_tensors[group_k] = group_v.param_groups[0]["params"][0] - else: - group_v.param_groups[0]["params"][0] = nn.Parameter(torch.cat((group_v.param_groups[0]["params"][0], extension_tensor), dim=0).requires_grad_(True)) - optimizable_tensors[group_k] = group_v.param_groups[0]["params"][0] - - return optimizable_tensors - - def densification_postfix(self, new_xyz, new_features_dc, new_features_rest, new_opacities, new_scaling, new_rotation): - d = {"xyz": new_xyz, - "f_dc": new_features_dc, - "f_rest": new_features_rest, - "opacity": new_opacities, - "scaling" : new_scaling, - "rotation" : new_rotation} - - optimizable_tensors = self.cat_tensors_to_optimizer(d) - self._xyz = optimizable_tensors["xyz"] - self._features_dc = optimizable_tensors["f_dc"] - self._features_rest = optimizable_tensors["f_rest"] - self._opacity = optimizable_tensors["opacity"] - self._scaling = optimizable_tensors["scaling"] - self._rotation = optimizable_tensors["rotation"] - - self.xyz_gradient_accum = torch.zeros((self.get_xyz.shape[0], 1), device="cuda") - self.denom = torch.zeros((self.get_xyz.shape[0], 1), device="cuda") - self.max_radii2D = torch.zeros((self.get_xyz.shape[0]), device="cuda") - - def densify_and_split(self, grads, grad_threshold, scene_extent, N=2, num_selected=None): - n_init_points = self.get_xyz.shape[0] - # Extract points that satisfy the gradient condition - padded_grad = torch.zeros((n_init_points), device="cuda") - padded_grad[:grads.shape[0]] = grads.squeeze() - selected_pts_mask = torch.where(padded_grad >= grad_threshold, True, False) - selected_pts_mask = torch.logical_and(selected_pts_mask, - torch.max(self.get_scaling, dim=1).values > self.percent_dense*scene_extent) - if num_selected is not None and selected_pts_mask.sum() > 0: - num_selected = int(min(num_selected, selected_pts_mask.sum() - 1)) - choices = np.arange(0, len(selected_pts_mask))[selected_pts_mask.cpu().numpy()] - chosen = np.random.choice(choices, num_selected, replace=False) - selected_pts_indices = torch.from_numpy(chosen).cuda() - selected_pts_mask[...] = False - selected_pts_mask[selected_pts_indices] = True - - stds = self.get_scaling[selected_pts_mask].repeat(N,1) - means =torch.zeros((stds.size(0), 3),device="cuda") - samples = torch.normal(mean=means, std=stds) - rots = gaussian_utils.build_rotation(self._rotation[selected_pts_mask]).repeat(N,1,1) - new_xyz = torch.bmm(rots, samples.unsqueeze(-1)).squeeze(-1) + self.get_xyz[selected_pts_mask].repeat(N, 1) - new_scaling = self.scaling_inverse_activation(self.get_scaling[selected_pts_mask].repeat(N,1) / (0.8*N)) - new_rotation = self._rotation[selected_pts_mask].repeat(N,1) - new_features_dc = self._features_dc[selected_pts_mask].repeat(N,1,1) - new_features_rest = self._features_rest[selected_pts_mask].repeat(N,1,1) - new_opacity = self._opacity[selected_pts_mask].repeat(N,1) - - self.densification_postfix(new_xyz, new_features_dc, new_features_rest, new_opacity, new_scaling, new_rotation) - - prune_filter = torch.cat((selected_pts_mask, torch.zeros(N * selected_pts_mask.sum(), device="cuda", dtype=bool))) - self.prune_points(prune_filter) - - def densify_and_clone(self, grads, grad_threshold, scene_extent, num_selected=None): - # Extract points that satisfy the gradient condition - selected_pts_mask = torch.where(torch.norm(grads, dim=-1) >= grad_threshold, True, False) - selected_pts_mask = torch.logical_and(selected_pts_mask, - torch.max(self.get_scaling, dim=1).values <= self.percent_dense*scene_extent) - - if num_selected is not None and selected_pts_mask.sum() > 0: - num_selected = int(min(num_selected, selected_pts_mask.sum() - 1)) - choices = np.arange(0, len(selected_pts_mask))[selected_pts_mask.cpu().numpy()] - chosen = np.random.choice(choices, num_selected, replace=False) - selected_pts_indices = torch.from_numpy(chosen).cuda() - selected_pts_mask[...] = False - selected_pts_mask[selected_pts_indices] = True - - # if selected_pts_mask.sum() > 0: - # breakpoint() - new_xyz = self._xyz[selected_pts_mask] - new_features_dc = self._features_dc[selected_pts_mask] - new_features_rest = self._features_rest[selected_pts_mask] - new_opacities = self._opacity[selected_pts_mask] - new_scaling = self._scaling[selected_pts_mask] - new_rotation = self._rotation[selected_pts_mask] - self.densification_postfix(new_xyz, new_features_dc, new_features_rest, new_opacities, new_scaling, new_rotation) - # if selected_pts_mask.sum() > 0: - # breakpoint() - - def densify_and_prune(self, max_grad, min_opacity, extent, max_screen_size): - grads = self.xyz_gradient_accum / self.denom - grads[grads.isnan()] = 0.0 - - num_added = self.max_points - self.get_xyz.shape[0] if self.max_points is not None else None - if num_added is None or num_added > 0: - self.densify_and_clone(grads, max_grad, extent, num_selected=num_added//2 if num_added is not None else None) - self.densify_and_split(grads, max_grad, extent, num_selected=num_added//2 if num_added is not None else None) - - prune_mask = (self.get_opacity < min_opacity).squeeze() - if max_screen_size: - big_points_vs = self.max_radii2D > max_screen_size - big_points_ws = self.get_scaling.max(dim=1).values > 0.1 * extent - prune_mask = torch.logical_or(torch.logical_or(prune_mask, big_points_vs), big_points_ws) - self.prune_points(prune_mask) - - torch.cuda.empty_cache() - - def add_densification_stats(self, viewspace_point_tensor, update_filter): - self.xyz_gradient_accum[update_filter] += torch.norm(viewspace_point_tensor.grad[update_filter,:2], dim=-1, keepdim=True) - self.denom[update_filter] += 1 - - def reset_opacity(self): - opacities_new = gaussian_utils.inverse_sigmoid(torch.min(self.get_opacity, torch.ones_like(self.get_opacity)*0.01)) - optimizable_tensors = self.replace_tensor_to_optimizer(opacities_new, "opacity") - self._opacity = optimizable_tensors["opacity"] diff --git a/nerfstudio/models/gaussian_splatting_inhouse.py b/nerfstudio/models/gaussian_splatting_inhouse.py deleted file mode 100644 index fff11b4f8..000000000 --- a/nerfstudio/models/gaussian_splatting_inhouse.py +++ /dev/null @@ -1,704 +0,0 @@ -# Copyright 2022 the Regents of the University of California, Nerfstudio Team and contributors. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -""" -NeRF implementation that combines many recent advancements. -""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from typing import Any, Dict, List, Literal, Optional, Tuple, Type, Union -import torch -from nerfstudio.data.scene_box import OrientedBox -from copy import deepcopy -from nerfstudio.cameras.rays import RayBundle -from torch.nn import Parameter -from torchmetrics.functional import structural_similarity_index_measure -from torchmetrics.image import PeakSignalNoiseRatio -from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity -from diff_rast._torch_impl import compute_sh_color -from nerfstudio.cameras.cameras import Cameras -from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes, TrainingCallbackLocation -from nerfstudio.viewer_beta.viewer_elements import ViewerControl -from nerfstudio.engine.optimizers import Optimizers -from nerfstudio.models.base_model import Model, ModelConfig -import math -import numpy as np -from sklearn.neighbors import NearestNeighbors - - -from diff_rast.rasterize import RasterizeGaussians -from diff_rast.project_gaussians import ProjectGaussians -from diff_rast.sh import SphericalHarmonics, num_sh_bases - -def random_quat_tensor(N, **kwargs): - u = torch.rand(N, **kwargs) - v = torch.rand(N, **kwargs) - w = torch.rand(N, **kwargs) - return torch.stack( - [ - torch.sqrt(1 - u) * torch.sin(2 * math.pi * v), - torch.sqrt(1 - u) * torch.cos(2 * math.pi * v), - torch.sqrt(u) * torch.sin(2 * math.pi * w), - torch.sqrt(u) * torch.sin(2 * math.pi * w), - ], - dim=-1, - ) - -def projection_matrix(znear, zfar, fovx, fovy, **kwargs): - t = znear * math.tan(0.5 * fovy) - b = -t - r = znear * math.tan(0.5 * fovx) - l = -r - n = znear - f = zfar - return torch.tensor( - [ - [2 * n / (r - l), 0.0, (r + l) / (r - l), 0.0], - [0.0, 2 * n / (t - b), (t + b) / (t - b), 0.0], - [0.0, 0.0, (f + n) / (f - n), -1.0 * f * n / (f - n)], - [0.0, 0.0, 1.0, 0.0], - ], - **kwargs, - ) - - -@dataclass -class GaussianSplattingModelConfig(ModelConfig): - """Gaussian Splatting Model Config""" - - _target: Type = field(default_factory=lambda: GaussianSplattingModel) - warmup_length:int = 1000 - """period of steps where refinement is turned off""" - refine_every: int = 100 - """period of steps where gaussians are culled and densified""" - resolution_schedule: int = 200 - """training starts at 1/d resolution, every n steps this is doubled""" - num_downscales: int = 2 - """at the beginning, resolution is 1/2^d, where d is this number""" - cull_alpha_thresh: float = 0.02 - """threshold of opacity for culling gaussians""" - cull_scale_thresh: float = 0.5 - """threshold of scale for culling gaussians""" - reset_alpha_every: int = 30 - """Every this many refinement steps, reset the alpha""" - densify_grad_thresh: float = 0.0002 - """threshold of positional gradient norm for densifying gaussians""" - densify_size_thresh: float = 0.01 - """below this size, gaussians are *duplicated*, otherwise split""" - sh_degree_interval: int = 1000 - """every n intervals turn on another sh degree""" - max_screen_size: int = 300 - """maximum screen size of a gaussian, in pixels""" - random_init: bool = False - """whether to initialize the positions uniformly randomly (not SFM points)""" - ssim_lambda: float = .2 - """weight of ssim loss""" - stop_split_at: int = 25000 - """stop splitting at this step""" - sh_degree: int = 4 - """maximum degree of spherical harmonics to use""" - -class GaussianSplattingModel(Model): - """Gaussian Splatting model - - Args: - config: Gaussian Splatting configuration to instantiate model - """ - - config: GaussianSplattingModelConfig - - def __init__(self, *args, **kwargs): - if "seed_points" in kwargs: - self.seed_pts = kwargs["seed_points"] - else: - self.seed_pts = None - super().__init__(*args, **kwargs) - self.vc = ViewerControl() - - def populate_modules(self): - if self.seed_pts is not None and not self.config.random_init: - self.means = torch.nn.Parameter(self.seed_pts[0]) # (Location, Color) - else: - self.means = torch.nn.Parameter((torch.rand((100000, 3)) - 0.5) * 2) - self.xys_grad_norm = None - self.max_2Dsize = None - distances, _ = self.k_nearest_sklearn(self.means.data, 3) - distances = torch.from_numpy(distances) - # find the average of the three nearest neighbors for each point and use that as the scale - avg_dist = distances.mean(dim=-1, keepdim=True) - self.scales = torch.nn.Parameter(torch.log(avg_dist.repeat(1, 3))) - - self.quats = torch.nn.Parameter(random_quat_tensor(self.num_points)) - dim_sh = num_sh_bases(self.config.sh_degree) - - if self.seed_pts is not None and not self.config.random_init: - fused_color = self.RGB2SH(self.seed_pts[1] / 255) - shs = torch.zeros((fused_color.shape[0], 3, dim_sh)).float().cuda() - shs[:, :3, 0] = fused_color - shs[:, 3:, 1:] = 0.0 - self.colors = torch.nn.Parameter(shs[:, :, 0:1]) - self.shs_rest = torch.nn.Parameter(shs[:, :, 1:]) - else: - self.colors = torch.nn.Parameter(torch.rand(self.num_points, 3, 1)) - self.shs_rest = torch.nn.Parameter(torch.zeros((self.num_points, 3, dim_sh - 1))) - - self.opacities = torch.nn.Parameter(torch.logit(0.1 * torch.ones(self.num_points, 1))) - # metrics - self.psnr = PeakSignalNoiseRatio(data_range=1.0) - from torchmetrics.image import StructuralSimilarityIndexMeasure - - self.ssim = StructuralSimilarityIndexMeasure() - self.lpips = LearnedPerceptualImagePatchSimilarity(normalize=True) - self.step = 0 - self.crop_box:Optional[OrientedBox] = None - self.back_color = torch.zeros(3) - - @property - def get_colors(self): - color = self.colors - shs_rest = self.shs_rest - return torch.cat((color, shs_rest), dim=2) - - def load_state_dict(self, dict, **kwargs): - # resize the parameters to match the new number of points - self.step = 30000 - newp = dict['means'].shape[0] - self.means = torch.nn.Parameter(torch.zeros(newp,3,device=self.device)) - self.scales = torch.nn.Parameter(torch.zeros(newp,3,device=self.device)) - self.quats = torch.nn.Parameter(torch.zeros(newp,4,device=self.device)) - self.colors = torch.nn.Parameter(torch.zeros(self.num_points, 3, 1,device=self.device)) - self.shs_rest = torch.nn.Parameter(torch.zeros(self.num_points, 3, num_sh_bases(self.config.sh_degree)-1,device=self.device)) - self.opacities = torch.nn.Parameter(torch.zeros(newp,1,device=self.device)) - super().load_state_dict(dict,**kwargs) - - def k_nearest_sklearn(self, x: torch.Tensor, k: int): - """ - Find k-nearest neighbors using sklearn's NearestNeighbors. - x: The data tensor of shape [num_samples, num_features] - k: The number of neighbors to retrieve - """ - # Convert tensor to numpy array - x_np = x.cpu().numpy() - - # Build the nearest neighbors model - nn_model = NearestNeighbors(n_neighbors=k + 1, algorithm="auto", metric="euclidean").fit(x_np) - - # Find the k-nearest neighbors - distances, indices = nn_model.kneighbors(x_np) - - # Exclude the point itself from the result and return - return distances[:, 1:].astype(np.float32), indices[:, 1:].astype(np.float32) - - def remove_from_optim(self, optimizer, deleted_mask, new_params): - """removes the deleted_mask from the optimizer provided""" - assert len(new_params) == 1 - assert isinstance(optimizer, torch.optim.Adam), "Only works with Adam" - - param = optimizer.param_groups[0]["params"][0] - param_state = optimizer.state[param] - del optimizer.state[param] - - # Modify the state directly without deleting and reassigning. - param_state['exp_avg'] = param_state['exp_avg'][~deleted_mask] - param_state['exp_avg_sq'] = param_state['exp_avg_sq'][~deleted_mask] - - # Update the parameter in the optimizer's param group. - del optimizer.param_groups[0]["params"][0] - del optimizer.param_groups[0]["params"] - optimizer.param_groups[0]["params"] = new_params - optimizer.state[new_params[0]] = param_state - - def dup_in_optim(self,optimizer, dup_mask, new_params,n=2): - """adds the parameters to the optimizer""" - param = optimizer.param_groups[0]["params"][0] - param_state = optimizer.state[param] - repeat_dims = (n,) + tuple(1 for _ in range(param_state['exp_avg'].dim() - 1)) - param_state['exp_avg'] = torch.cat([param_state['exp_avg'], torch.zeros_like(param_state['exp_avg'][dup_mask.squeeze()]).repeat(*repeat_dims)],dim=0) - param_state['exp_avg_sq'] = torch.cat([param_state['exp_avg_sq'], torch.zeros_like(param_state['exp_avg_sq'][dup_mask.squeeze()]).repeat(*repeat_dims)],dim=0) - del optimizer.state[param] - optimizer.state[new_params[0]] = param_state - optimizer.param_groups[0]["params"] = new_params - del param - - def after_train(self, step): - with torch.no_grad(): - # keep track of a moving average of grad norms - visible_mask = (self.radii > 0).flatten() - grads = self.xys.grad.detach().norm(dim=-1)#TODO fill in - # print(f"grad norm min {grads.min().item()} max {grads.max().item()} mean {grads.mean().item()} size {grads.shape}") - if self.xys_grad_norm is None: - self.xys_grad_norm = grads - self.vis_counts = torch.ones_like(self.xys_grad_norm) - else: - self.vis_counts[visible_mask] = self.vis_counts[visible_mask] + 1 - self.xys_grad_norm[visible_mask] = grads[visible_mask] + self.xys_grad_norm[visible_mask] - - # update the max screen size - if self.max_2Dsize is None: - self.max_2Dsize = torch.zeros_like(self.radii) - newradii = self.radii[visible_mask] - self.max_2Dsize[visible_mask] = torch.maximum(self.max_2Dsize[visible_mask], newradii) - - def set_crop(self,crop_box:OrientedBox): - self.crop_box = crop_box - - def refinement_after(self, optimizers: Optimizers, step): - if self.step > self.config.warmup_length: - with torch.no_grad(): - if self.step self.config.densify_grad_thresh).squeeze() - splits = (self.scales.exp().max(dim=-1).values > self.config.densify_size_thresh).squeeze() - splits &= high_grads - nsamps=2 - split_means,split_colors,split_shs,split_opacities,split_scales,split_quats = self.split_gaussians(splits,nsamps) - - dups = (self.scales.exp().max(dim=-1).values <= self.config.densify_size_thresh).squeeze() - dups &= high_grads - dup_means,dup_colors,dup_shs,dup_opacities,dup_scales,dup_quats = self.dup_gaussians(dups) - self.means = Parameter(torch.cat([self.means.detach(),split_means, dup_means],dim=0)) - self.colors = Parameter(torch.cat([self.colors.detach(),split_colors, dup_colors],dim=0)) - self.shs_rest = Parameter(torch.cat([self.shs_rest.detach(),split_shs, dup_shs],dim=0)) - self.opacities = Parameter(torch.cat([self.opacities.detach(),split_opacities, dup_opacities],dim=0)) - self.scales = Parameter(torch.cat([self.scales.detach(),split_scales, dup_scales],dim=0)) - self.quats = Parameter(torch.cat([self.quats.detach(),split_quats, dup_quats],dim=0)) - #append zeros to the max_2Dsize tensor - self.max_2Dsize = torch.cat([self.max_2Dsize,torch.zeros_like(split_scales[:,0]),torch.zeros_like(dup_scales[:,0])],dim=0) - split_idcs = torch.where(splits)[0] - param_groups = self.get_param_groups() - for group,param in param_groups.items(): - self.dup_in_optim(optimizers.optimizers[group],split_idcs,param,n=nsamps) - dup_idcs = torch.where(dups)[0] - - param_groups = self.get_param_groups() - for group,param in param_groups.items(): - self.dup_in_optim(optimizers.optimizers[group],dup_idcs,param,1) - - #only cull if we've seen every image since opacity reset - reset_interval = (self.config.reset_alpha_every * self.config.refine_every) - if self.step % reset_interval > self.num_train_data: - #then cull - deleted_mask = self.cull_gaussians() - param_groups = self.get_param_groups() - for group, param in param_groups.items(): - self.remove_from_optim(optimizers.optimizers[group], deleted_mask, param) - - if self.step // self.config.refine_every % self.config.reset_alpha_every == 0: - reset_value = 0.01 - self.opacities.data = torch.full_like( - self.opacities.data, torch.logit(torch.tensor(reset_value)).item() - ) - # reset the exp of optimizer - optim = optimizers.optimizers["opacity"] - param = optim.param_groups[0]["params"][0] - param_state = optim.state[param] - param_state["exp_avg"] = torch.zeros_like(param_state["exp_avg"]) - param_state["exp_avg_sq"] = torch.zeros_like(param_state["exp_avg_sq"]) - self.xys_grad_norm = None - self.vis_counts = None - self.max_2Dsize = None - - def cull_gaussians(self): - """ - This function deletes gaussians with under a certain opacity threshold - """ - n_bef = self.num_points - # cull transparent ones - culls = (torch.sigmoid(self.opacities) < self.config.cull_alpha_thresh).squeeze() - if self.step > self.config.refine_every * self.config.reset_alpha_every: - # cull huge ones - culls = culls | (torch.exp(self.scales).max(dim=-1).values > self.config.cull_scale_thresh).squeeze() - # cull big screen space - culls = culls | (self.max_2Dsize > self.config.max_screen_size).squeeze() - self.means = Parameter(self.means[~culls].detach()) - self.scales = Parameter(self.scales[~culls].detach()) - self.quats = Parameter(self.quats[~culls].detach()) - self.colors = Parameter(self.colors[~culls].detach()) - self.shs_rest = Parameter(self.shs_rest[~culls].detach()) - self.opacities = Parameter(self.opacities[~culls].detach()) - - print(f"Culled {n_bef - self.num_points} gaussians") - return culls - - def split_gaussians(self, split_mask,samps): - """ - This function splits gaussians that are too large - """ - from diff_rast._torch_impl import quat_to_rotmat - n_splits = split_mask.sum().item() - print(f"Splitting {split_mask.sum().item()/self.num_points} gaussians: {n_splits}/{self.num_points}") - centered_samples = torch.randn((samps*n_splits,3),device=self.device) # Nx3 of axis-aligned scales - scaled_samples = torch.exp(self.scales[split_mask].repeat(samps,1)) * centered_samples # how these scales are rotated - quats = self.quats[split_mask]/self.quats[split_mask].norm(dim=-1,keepdim=True) #normalize them first - rots = quat_to_rotmat(quats.repeat(samps,1)) # how these scales are rotated - rotated_samples = torch.bmm(rots,scaled_samples[...,None]).squeeze() - new_means = rotated_samples + self.means[split_mask].repeat(samps,1) - # step 2, sample new colors - new_colors = self.colors[split_mask].repeat(samps,1,1) - # step 3, sample new shs - new_shs = self.shs_rest[split_mask].repeat(samps,1,1) - # step 3, sample new opacities - new_opacities = self.opacities[split_mask].repeat(samps,1) - # step 4, sample new scales - new_scales = torch.log(torch.exp(self.scales[split_mask])/1.6).repeat(samps,1) - self.scales[split_mask] = torch.log(torch.exp(self.scales[split_mask])/1.6) - # step 5, sample new quats - new_quats = self.quats[split_mask].repeat(samps,1) - return new_means,new_colors,new_shs,new_opacities,new_scales,new_quats - - def dup_gaussians(self, dup_mask): - """ - This function duplicates gaussians that are too small - """ - n_dups = dup_mask.sum().item() - print(f"Would duplicate {dup_mask.sum().item()/self.num_points} gaussians: {n_dups}/{self.num_points}") - dup_means = self.means[dup_mask] - dup_colors = self.colors[dup_mask] - dup_shs = self.shs_rest[dup_mask] - dup_opacities = self.opacities[dup_mask] - dup_scales = self.scales[dup_mask] - dup_quats = self.quats[dup_mask] - return dup_means, dup_colors, dup_shs, dup_opacities, dup_scales, dup_quats - - @property - def num_points(self): - return self.means.shape[0] - - def get_training_callbacks( - self, training_callback_attributes: TrainingCallbackAttributes - ) -> List[TrainingCallback]: - cbs = [] - cbs.append(TrainingCallback([TrainingCallbackLocation.BEFORE_TRAIN_ITERATION], self.step_cb)) - # The order of these matters - cbs.append( - TrainingCallback( - [TrainingCallbackLocation.AFTER_TRAIN_ITERATION], - self.after_train, - ) - ) - cbs.append( - TrainingCallback( - [TrainingCallbackLocation.AFTER_TRAIN_ITERATION], - self.refinement_after, - update_every_num_iters=self.config.refine_every, - args=[training_callback_attributes.optimizers], - ) - ) - return cbs - - def step_cb(self, step): - self.step = step - - def get_param_groups(self) -> Dict[str, List[Parameter]]: - """Obtain the parameter groups for the optimizers - - Returns: - Mapping of different parameter groups - """ - return { - "xyz": [self.means], - "color": [self.colors], - "shs": [self.shs_rest], - "opacity": [self.opacities], - "scaling": [self.scales], - "rotation": [self.quats], - } - - def _get_downscale_factor(self): - return 2 ** max((self.config.num_downscales - self.step // self.config.resolution_schedule), 0) - - def get_outputs(self, camera: Cameras) -> Dict[str, Union[torch.Tensor, List]]: - """Takes in a Ray Bundle and returns a dictionary of outputs. - - Args: - ray_bundle: Input bundle of rays. This raybundle should have all the - needed information to compute the outputs. - - Returns: - Outputs of model. (ie. rendered colors) - """ - import viser.transforms as vtf - import numpy as np - - if not isinstance(camera, Cameras): - print("Called get_outputs with not a camera") - return {} - assert camera.shape[0] == 1, "Only one camera at a time" - # dont mutate the input - camera = deepcopy(camera) - if self.training: - d = self._get_downscale_factor() - camera.rescale_output_resolution(1/d) - if self.crop_box is not None and not self.training: - crop_ids = self.crop_box.within(self.means).squeeze() - if crop_ids.sum()==0: - return {'rgb':torch.full((camera.height.item(),camera.width.item(),3),.5,device=self.device)} - else: - crop_ids = torch.ones_like(self.means[:,0],dtype=torch.bool) - #shift the camera to center of scene looking at center - R = camera.camera_to_worlds[..., :3, :3] # 1 x 3 x 3 - T = camera.camera_to_worlds[..., :3, 3:4] # 1 x 3 x 1 - R = vtf.SO3.from_matrix(R.cpu().squeeze().numpy()) - R = R @ vtf.SO3.from_x_radians(np.pi) - R = torch.from_numpy(R.as_matrix()[None, ...]).to(self.device, torch.float32) - viewmat = torch.cat([R, T], dim=2) - # add a row of zeros and a 1 to the bottom of the viewmat - viewmat = torch.cat([viewmat, torch.tensor([[[0, 0, 0, 1]]], device=self.device)], dim=1) - # invert it - viewmat = torch.inverse(viewmat) - # calculate the FOV of the camera given fx and fy, width and height - cx = camera.cx.item() - cy = camera.cy.item() - fovx = 2 * math.atan(camera.width / (2 * camera.fx)) - fovy = 2 * math.atan(camera.height / (2 * camera.fy)) - W, H = camera.width.item(), camera.height.item() - self.last_size = (H, W) - projmat = projection_matrix(0.001, 1000, fovx, fovy).to(self.device) - BLOCK_X, BLOCK_Y = 16, 16 - tile_bounds = ( - (W + BLOCK_X - 1) // BLOCK_X, - (H + BLOCK_Y - 1) // BLOCK_Y, - 1, - ) - if self.training: - background = torch.rand(3, device=self.device) - else: - background = torch.zeros(3, device=self.device) - self.dummy_2d_grads = torch.empty((self.num_points, 2), device=self.device) - self.xys, depths, self.radii, conics, num_tiles_hit, cov3d = ProjectGaussians.apply( - self.means[crop_ids], - self.dummy_2d_grads, - torch.exp(self.scales[crop_ids]), - 1, - self.quats[crop_ids]/self.quats[crop_ids].norm(dim=-1,keepdim=True), - viewmat.squeeze()[:3, :], - projmat.squeeze() @ viewmat.squeeze(), - camera.fx.item(), - camera.fy.item(), - cx, - cy, - H, - W, - tile_bounds, - ) - #Important to allow xys grads to populate properly - if self.training: - self.xys.retain_grad() - if self.config.sh_degree > 0: - viewdirs = self.means - camera.camera_to_worlds[..., :3, 3] # (N, 3) - viewdirs = viewdirs / viewdirs.norm(dim=-1, keepdim=True) - n = min(self.step // self.config.sh_degree_interval,self.config.sh_degree) - rgbs = eval_sh(n, self.get_colors, viewdirs)[crop_ids] - rgbs = torch.clamp(rgbs + 0.5,0.0,1.0) - else: - rgbs = self.get_colors.squeeze()[crop_ids] # (N, 3) - rgbs = torch.sigmoid(rgbs) - - rgb = RasterizeGaussians.apply( - self.xys, - depths, - self.radii, - conics, - num_tiles_hit, - rgbs, - torch.sigmoid(self.opacities[crop_ids]), - H, - W, - background, - ) - depth_im = RasterizeGaussians.apply( - self.xys, - depths, - self.radii, - conics, - num_tiles_hit, - depths[:,None], - torch.sigmoid(self.opacities[crop_ids]), - H, - W, - torch.ones(1,device=self.device)*10, - )[...,0:1] - return {"rgb": rgb[..., :3],'depth':depth_im} - - def get_metrics_dict(self, outputs, batch) -> Dict[str, torch.Tensor]: - """Compute and returns metrics. - - Args: - outputs: the output to compute loss dict to - batch: ground truth batch corresponding to outputs - """ - - return {} - - def get_loss_dict(self, outputs, batch, metrics_dict=None) -> Dict[str, torch.Tensor]: - """Computes and returns the losses dict. - - Args: - outputs: the output to compute loss dict to - batch: ground truth batch corresponding to outputs - metrics_dict: dictionary of metrics, some of which we can use for loss - """ - d = self._get_downscale_factor() - if d > 1: - # use torchvision to resize - import torchvision.transforms.functional as TF - - newsize = (batch["image"].shape[0] // d, batch["image"].shape[1] // d) - gt_img = TF.resize(batch["image"].permute(2, 0, 1), newsize).permute(1, 2, 0) - else: - gt_img = batch['image'] - Ll1 = torch.abs(gt_img- outputs['rgb']).mean() - simloss = (1-self.ssim(gt_img.permute(2,0,1)[None,...], outputs['rgb'].permute(2,0,1)[None,...])) - return {"main_loss": (1-self.config.ssim_lambda)*Ll1 + self.config.ssim_lambda*simloss} - - @torch.no_grad() - def get_outputs_for_camera_ray_bundle( - self, camera_ray_bundle: RayBundle, camera: Optional[Cameras] = None - ) -> Dict[str, torch.Tensor]: - """Takes in camera parameters and computes the output of the model. - - Args: - camera_ray_bundle: ray bundle to calculate outputs over - """ - assert camera is not None, "must provide camera to gaussian model" - outs = self.get_outputs(camera.to(self.device)) - return outs - - def get_image_metrics_and_images( - self, outputs: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor] - ) -> Tuple[Dict[str, float], Dict[str, torch.Tensor]]: - """Writes the test image outputs. - - Args: - image_idx: Index of the image. - step: Current step. - batch: Batch of data. - outputs: Outputs of the model. - - Returns: - A dictionary of metrics. - """ - gt_rgb = batch["image"].to(self.device) - predicted_rgb = outputs["rgb"] # Blended with background (black if random background) - - combined_rgb = torch.cat([gt_rgb, predicted_rgb], dim=1) - - # Switch images from [H, W, C] to [1, C, H, W] for metrics computations - gt_rgb = torch.moveaxis(gt_rgb, -1, 0)[None, ...] - predicted_rgb = torch.moveaxis(predicted_rgb, -1, 0)[None, ...] - - psnr = self.psnr(gt_rgb, predicted_rgb) - ssim = self.ssim(gt_rgb, predicted_rgb) - lpips = self.lpips(gt_rgb, predicted_rgb) - - # all of these metrics will be logged as scalars - metrics_dict = {"psnr": float(psnr.item()), "ssim": float(ssim)} # type: ignore - metrics_dict["lpips"] = float(lpips) - - images_dict = {"img": combined_rgb} - - return metrics_dict, images_dict - - def RGB2SH(self, rgb): - C0 = 0.28209479177387814 - return (rgb - 0.5) / C0 - - -def eval_sh(deg, sh, dirs): - """ - Evaluate spherical harmonics at unit directions - using hardcoded SH polynomials. - Works with torch/np/jnp. - ... Can be 0 or more batch dimensions. - Args: - deg: int SH deg. Currently, 0-3 supported - sh: jnp.ndarray SH coeffs [..., C, (deg + 1) ** 2] - dirs: jnp.ndarray unit directions [..., 3] - Returns: - [..., C] - """ - C0 = 0.28209479177387814 - C1 = 0.4886025119029199 - C2 = [1.0925484305920792, -1.0925484305920792, 0.31539156525252005, -1.0925484305920792, 0.5462742152960396] - C3 = [ - -0.5900435899266435, - 2.890611442640554, - -0.4570457994644658, - 0.3731763325901154, - -0.4570457994644658, - 1.445305721320277, - -0.5900435899266435, - ] - C4 = [ - 2.5033429417967046, - -1.7701307697799304, - 0.9461746957575601, - -0.6690465435572892, - 0.10578554691520431, - -0.6690465435572892, - 0.47308734787878004, - -1.7701307697799304, - 0.6258357354491761, - ] - - assert deg <= 4 and deg >= 0 - coeff = (deg + 1) ** 2 - assert sh.shape[-1] >= coeff - - result = C0 * sh[..., 0] - if deg > 0: - x, y, z = dirs[..., 0:1], dirs[..., 1:2], dirs[..., 2:3] - result = result - C1 * y * sh[..., 1] + C1 * z * sh[..., 2] - C1 * x * sh[..., 3] - - if deg > 1: - xx, yy, zz = x * x, y * y, z * z - xy, yz, xz = x * y, y * z, x * z - result = ( - result - + C2[0] * xy * sh[..., 4] - + C2[1] * yz * sh[..., 5] - + C2[2] * (2.0 * zz - xx - yy) * sh[..., 6] - + C2[3] * xz * sh[..., 7] - + C2[4] * (xx - yy) * sh[..., 8] - ) - - if deg > 2: - result = ( - result - + C3[0] * y * (3 * xx - yy) * sh[..., 9] - + C3[1] * xy * z * sh[..., 10] - + C3[2] * y * (4 * zz - xx - yy) * sh[..., 11] - + C3[3] * z * (2 * zz - 3 * xx - 3 * yy) * sh[..., 12] - + C3[4] * x * (4 * zz - xx - yy) * sh[..., 13] - + C3[5] * z * (xx - yy) * sh[..., 14] - + C3[6] * x * (xx - 3 * yy) * sh[..., 15] - ) - - if deg > 3: - result = ( - result - + C4[0] * xy * (xx - yy) * sh[..., 16] - + C4[1] * yz * (3 * xx - yy) * sh[..., 17] - + C4[2] * xy * (7 * zz - 1) * sh[..., 18] - + C4[3] * yz * (7 * zz - 3) * sh[..., 19] - + C4[4] * (zz * (35 * zz - 30) + 3) * sh[..., 20] - + C4[5] * xz * (7 * zz - 3) * sh[..., 21] - + C4[6] * (xx - yy) * (7 * zz - 1) * sh[..., 22] - + C4[7] * xz * (xx - 3 * yy) * sh[..., 23] - + C4[8] * (xx * (xx - 3 * yy) - yy * (3 * xx - yy)) * sh[..., 24] - ) - return result diff --git a/nerfstudio/viewer_beta/control_panel.py b/nerfstudio/viewer_beta/control_panel.py index 07f556dd4..b2c1c8ebb 100644 --- a/nerfstudio/viewer_beta/control_panel.py +++ b/nerfstudio/viewer_beta/control_panel.py @@ -270,10 +270,7 @@ class ControlPanel: step: the train step to set the model to """ with self.viser_server.atomic(), self.stat_folder: - # TODO change to a .value call instead of remove() and add, this makes it jittery - with self.viser_server.atomic(): - self.markdown.remove() - self.markdown = self.viser_server.add_gui_markdown(f"Step: {step}") + self.markdown.content=f"Step: {step}" def update_output_options(self, new_options: List[str]): """ diff --git a/nerfstudio/viewer_beta/render_panel.py b/nerfstudio/viewer_beta/render_panel.py index 7342b0f38..690027f8d 100644 --- a/nerfstudio/viewer_beta/render_panel.py +++ b/nerfstudio/viewer_beta/render_panel.py @@ -625,6 +625,7 @@ def populate_render_tab( # now write the json file json_outfile = datapath / "camera_paths" / f"{render_name_text.value}.json" + json_outfile.parent.mkdir(parents=True, exist_ok=True) with open(json_outfile.absolute(), "w") as outfile: json.dump(json_data, outfile) # now show the command