From 4f462ea7375d96e183b8558ae5f0ae0fea0b07e7 Mon Sep 17 00:00:00 2001 From: Justin Kerr Date: Sat, 30 Sep 2023 00:27:24 -0700 Subject: [PATCH 1/2] fix order of cull and split, add random init support --- nerfstudio/configs/method_configs.py | 2 +- .../models/gaussian_splatting_inhouse.py | 71 +++++++++---------- 2 files changed, 35 insertions(+), 38 deletions(-) diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index a12232c85..5f057e766 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -594,7 +594,7 @@ method_configs["gaussian-splatting"] = TrainerConfig( steps_per_eval_all_images=1000000, # set to a very large model so we don't eval with all images max_num_iterations=3000000, mixed_precision=False, - gradient_accumulation_steps=4, + gradient_accumulation_steps=5, pipeline=VanillaPipelineConfig( datamanager=FullImageDatamanagerConfig( dataparser=ColmapDataParserConfig(load_3D_points=True), diff --git a/nerfstudio/models/gaussian_splatting_inhouse.py b/nerfstudio/models/gaussian_splatting_inhouse.py index 021350e70..a1296caae 100644 --- a/nerfstudio/models/gaussian_splatting_inhouse.py +++ b/nerfstudio/models/gaussian_splatting_inhouse.py @@ -88,11 +88,11 @@ def projection_matrix(znear, zfar, fovx, fovy, **kwargs): class GaussianSplattingModelConfig(ModelConfig): """Gaussian Splatting Model Config""" _target: Type = field(default_factory=lambda: GaussianSplattingModel) - warmup_length:int = 1500 + warmup_length:int = 900 """period of steps where refinement is turned off""" refine_every:int = 100 """period of steps where gaussians are culled and densified""" - resolution_schedule: int = 500 + resolution_schedule: int = 300 """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""" @@ -107,10 +107,12 @@ class GaussianSplattingModelConfig(ModelConfig): densify_size_thresh: float = .01 """below this size, gaussians are *duplicated*, otherwise split""" #set to 3 to effeectively turn off, it seems to make things worse? need to investigate - sh_degree_interval: int = 3 + sh_degree_interval: int = 1000 """every n intervals turn on another sh degree""" - max_screen_size: int = 200 + max_screen_size: int = 1000 """maximum screen size of a gaussian, in pixels""" + random_init: bool = False + """whether to initialize the positions uniformly randomly (not SFM points)""" class GaussianSplattingModel(Model): """Gaussian Splatting model @@ -121,18 +123,24 @@ class GaussianSplattingModel(Model): config: GaussianSplattingModelConfig def __init__(self,*args,**kwargs): - self.seed_pts = kwargs['seed_points'] + 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): # TODO (jake-austin): clean this up, this is transplanted code across all the implementation functions - self.means = torch.nn.Parameter(self.seed_pts[0]) # (Location, Color) + 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))-.5)*2) self.means_grad_norm = None init_scale = torch.log(torch.tensor(.01)).item() self.scales = torch.nn.Parameter(torch.full((self.num_points,3),init_scale)) self.max_2Dsize = None - distances, _ = self.k_nearest_sklearn(self.seed_pts[0], 3) + 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) @@ -142,15 +150,17 @@ class GaussianSplattingModel(Model): self.quats = torch.nn.Parameter(random_quat_tensor(self.num_points)) self.degree = 4 dim_sh = num_sh_bases(self.degree) - # self.colors = torch.nn.Parameter(torch.rand(self.num_points, dim_sh, 3)) - fused_color = self.RGB2SH(self.seed_pts[1]/255) - shs = torch.zeros((fused_color.shape[0], dim_sh, 3)).float().cuda() - shs[:, 0, :] = fused_color - shs[:, 1:, :] = 0.0 - self.colors = torch.nn.Parameter(shs) + 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], dim_sh, 3)).float().cuda() + shs[:, 0, :] = fused_color + shs[:, 1:, :] = 0.0 + self.colors = torch.nn.Parameter(shs) + else: + self.colors = torch.nn.Parameter(torch.rand(self.num_points, dim_sh, 3)) - self.opacities = torch.nn.Parameter(torch.logit(0.1 * torch.ones(self.num_points, 1))) + self.opacities = torch.nn.Parameter(torch.logit(0.01 * torch.ones(self.num_points, 1))) # metrics self.psnr = PeakSignalNoiseRatio(data_range=1.0) from torchmetrics.image import StructuralSimilarityIndexMeasure @@ -246,19 +256,6 @@ class GaussianSplattingModel(Model): newradii = self.radii[visible_mask] self.max_2Dsize[visible_mask] = torch.maximum(self.max_2Dsize[visible_mask],newradii) - - def refinement_before(self, optimizers:Optimizers, step): - print("Inside refinement before") - if self.step > self.config.warmup_length: - with torch.no_grad(): - # do all the refinement stuff here - #first we cull gaussians - 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) - - def refinement_after(self, optimizers:Optimizers, step): if self.step > self.config.warmup_length: with torch.no_grad(): @@ -277,6 +274,8 @@ class GaussianSplattingModel(Model): 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] dup_idcs = torch.where(dups)[0] @@ -285,6 +284,12 @@ class GaussianSplattingModel(Model): param_groups = self.get_param_groups() for group,param in param_groups.items(): self.dup_in_optim(optimizers.optimizers[group],add_idcs,param) + #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: print("Resetting alpha") @@ -312,9 +317,6 @@ class GaussianSplattingModel(Model): #cull big screen space if self.step>self.config.refine_every*self.config.reset_alpha_every: culls = culls | (self.max_2Dsize > self.config.max_screen_size).squeeze() - self.means_grad_norm = self.means_grad_norm[~culls].detach() - self.max_2Dsize = self.max_2Dsize[~culls].detach() - self.vis_counts = self.vis_counts[~culls].detach() self.means = Parameter(self.means[~culls].detach()) self.scales = Parameter(self.scales[~culls].detach()) self.quats = Parameter(self.quats[~culls].detach()) @@ -361,7 +363,7 @@ class GaussianSplattingModel(Model): """ 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_means = self.means[dup_mask] + self.means.grad[dup_mask] dup_colors = self.colors[dup_mask] dup_opacities = self.opacities[dup_mask] dup_scales = self.scales[dup_mask] @@ -377,11 +379,6 @@ class GaussianSplattingModel(Model): ) -> List[TrainingCallback]: cbs = [] cbs.append(TrainingCallback([TrainingCallbackLocation.BEFORE_TRAIN_ITERATION],self.step_cb)) - cbs.append(TrainingCallback([TrainingCallbackLocation.BEFORE_TRAIN_ITERATION], - self.refinement_before, - update_every_num_iters=self.config.refine_every, - args=[training_callback_attributes.optimizers] - )) # The order of these matters cbs.append(TrainingCallback([TrainingCallbackLocation.AFTER_TRAIN_ITERATION], self.after_train,)) @@ -454,7 +451,7 @@ class GaussianSplattingModel(Model): 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() - projmat = projection_matrix(.0001,1000,fovx,fovy).to(self.device) + projmat = projection_matrix(.1,1000,fovx,fovy).to(self.device) BLOCK_X, BLOCK_Y = 16, 16 tile_bounds = ( (W + BLOCK_X - 1) // BLOCK_X, From b350a68ff4d6289074220044b1359b855eb722c5 Mon Sep 17 00:00:00 2001 From: Justin Kerr Date: Sat, 30 Sep 2023 11:01:45 -0700 Subject: [PATCH 2/2] integrate render pipeline (WIP), fix ssim loss --- nerfstudio/configs/method_configs.py | 6 +- nerfstudio/models/base_model.py | 3 +- .../models/gaussian_splatting_inhouse.py | 57 ++++++++----------- nerfstudio/scripts/render.py | 4 +- nerfstudio/scripts/viewer/run_viewer.py | 2 +- .../viewer_beta/render_state_machine.py | 4 +- 6 files changed, 34 insertions(+), 42 deletions(-) diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 5f057e766..9bf695d9c 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -590,11 +590,11 @@ method_configs["gaussian-splatting"] = TrainerConfig( method_name="gaussian-splatting", steps_per_eval_image=10, steps_per_eval_batch=10, - steps_per_save=2000000,#saving is broken + steps_per_save=2000, steps_per_eval_all_images=1000000, # set to a very large model so we don't eval with all images - max_num_iterations=3000000, + max_num_iterations=30000, mixed_precision=False, - gradient_accumulation_steps=5, + gradient_accumulation_steps=1, pipeline=VanillaPipelineConfig( datamanager=FullImageDatamanagerConfig( dataparser=ColmapDataParserConfig(load_3D_points=True), diff --git a/nerfstudio/models/base_model.py b/nerfstudio/models/base_model.py index aa5562895..9f0ae77d9 100644 --- a/nerfstudio/models/base_model.py +++ b/nerfstudio/models/base_model.py @@ -164,7 +164,8 @@ class Model(nn.Module): """ @torch.no_grad() - def get_outputs_for_camera_ray_bundle(self, camera_ray_bundle: Union[RayBundle, Cameras]) -> Dict[str, torch.Tensor]: + def get_outputs_for_camera_ray_bundle(self, camera_ray_bundle: Union[RayBundle, Cameras], + camera:Optional[Cameras] = None) -> Dict[str, torch.Tensor]: """Takes in camera parameters and computes the output of the model. TODO (jake-austin): Make 2 cases for when this is a ray or camera based model Args: diff --git a/nerfstudio/models/gaussian_splatting_inhouse.py b/nerfstudio/models/gaussian_splatting_inhouse.py index a1296caae..cb705eeff 100644 --- a/nerfstudio/models/gaussian_splatting_inhouse.py +++ b/nerfstudio/models/gaussian_splatting_inhouse.py @@ -22,6 +22,7 @@ from dataclasses import dataclass, field from typing import Any, Dict, List, Literal, Optional, Tuple, Type, Union import torch 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 @@ -109,10 +110,11 @@ class GaussianSplattingModelConfig(ModelConfig): #set to 3 to effeectively turn off, it seems to make things worse? need to investigate sh_degree_interval: int = 1000 """every n intervals turn on another sh degree""" - max_screen_size: int = 1000 + max_screen_size: int = 700 """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 class GaussianSplattingModel(Model): """Gaussian Splatting model @@ -133,7 +135,8 @@ class GaussianSplattingModel(Model): def populate_modules(self): # TODO (jake-austin): clean this up, this is transplanted code across all the implementation functions if self.seed_pts is not None and not self.config.random_init: - self.means = torch.nn.Parameter(self.seed_pts[0]) # (Location, Color) + randmeans = torch.rand((30000,3))*8-4 + self.means = torch.nn.Parameter(torch.cat([self.seed_pts[0],randmeans],dim=0)) # (Location, Color) else: self.means = torch.nn.Parameter((torch.rand((100000,3))-.5)*2) self.means_grad_norm = None @@ -153,10 +156,11 @@ class GaussianSplattingModel(Model): 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], dim_sh, 3)).float().cuda() + shs = torch.zeros((fused_color.shape[0], dim_sh, 3)) shs[:, 0, :] = fused_color shs[:, 1:, :] = 0.0 - self.colors = torch.nn.Parameter(shs) + randcolors = torch.rand((30000,dim_sh,3)) + self.colors = torch.nn.Parameter(torch.cat([shs,randcolors])) else: self.colors = torch.nn.Parameter(torch.rand(self.num_points, dim_sh, 3)) @@ -167,34 +171,24 @@ class GaussianSplattingModel(Model): self.ssim = StructuralSimilarityIndexMeasure() self.lpips = LearnedPerceptualImagePatchSimilarity(normalize=True) self.step=0 - - def k_nearest_faiss(self, x, k): - """ - Find k-nearest neighbors using FAISS. - x: The data matrix (NumPy array) of shape [num_samples, num_features] - k: The number of neighbors to retrieve - Returns: A tuple (distances, indices) where both are Numpy arrays of shape [num_samples, k] - """ - num_samples = x.shape[0] - - # Build the index - index = faiss.IndexFlatL2(x.shape[1]) # Here, using a flat (brute force) index for L2 distance - index.add(x) - - # Search the index - distances, indices = index.search(x, k + 1) # +1 because the point itself is always returned - - # Exclude the point itself from the result and return - return distances[:, 1:], indices[:, 1:] + 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, num_sh_bases(self.degree), 3,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 """ - num_samples = x.size(0) - # Convert tensor to numpy array x_np = x.cpu().numpy() @@ -206,7 +200,6 @@ class GaussianSplattingModel(Model): # 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""" @@ -227,7 +220,6 @@ class GaussianSplattingModel(Model): optimizer.param_groups[0]['params'] = new_params optimizer.state[new_params[0]] = param_state - def dup_in_optim(self,optimizer, dup_mask, new_params): """adds the parameters to the optimizer""" param = optimizer.param_groups[0]['params'][0] @@ -529,19 +521,18 @@ class GaussianSplattingModel(Model): else: gt_img = batch['image'] Ll1 = torch.nn.functional.l1_loss(gt_img, outputs['rgb']) - # This simloss makes the results look weird, removing for now - # simloss = self.ssim(gt_img.permute(2,0,1)[None,...], outputs['rgb'].permute(2,0,1)[None,...]) - return {"main_loss": Ll1} + simloss = (1-self.ssim(gt_img.permute(2,0,1)[None,...], outputs['rgb'].permute(2,0,1)[None,...]))/2.0 + 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( diff --git a/nerfstudio/scripts/render.py b/nerfstudio/scripts/render.py index 99a7cae01..65e98764d 100644 --- a/nerfstudio/scripts/render.py +++ b/nerfstudio/scripts/render.py @@ -136,10 +136,10 @@ def _render_trajectory_video( with renderers.background_color_override_context( crop_data.background_color.to(pipeline.device) ), torch.no_grad(): - outputs = pipeline.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle) + outputs = pipeline.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle,camera=cameras[camera_idx:camera_idx+1]) else: with torch.no_grad(): - outputs = pipeline.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle) + outputs = pipeline.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle,camera=cameras[camera_idx:camera_idx+1]) render_image = [] for rendered_output_name in rendered_output_names: diff --git a/nerfstudio/scripts/viewer/run_viewer.py b/nerfstudio/scripts/viewer/run_viewer.py index de210a30d..b1b87c42c 100644 --- a/nerfstudio/scripts/viewer/run_viewer.py +++ b/nerfstudio/scripts/viewer/run_viewer.py @@ -104,7 +104,7 @@ def _start_viewer(config: TrainerConfig, pipeline: Pipeline, step: int): viewer_state = ViewerBetaState( config.viewer, log_filename=viewer_log_path, - datapath=base_dir, + datapath=pipeline.datamanager.get_datapath(), pipeline=pipeline, share=config.viewer.make_share_url, ) diff --git a/nerfstudio/viewer_beta/render_state_machine.py b/nerfstudio/viewer_beta/render_state_machine.py index 07c32ea77..f75f91368 100644 --- a/nerfstudio/viewer_beta/render_state_machine.py +++ b/nerfstudio/viewer_beta/render_state_machine.py @@ -142,10 +142,10 @@ class RenderStateMachine(threading.Thread): with background_color_override_context( background_color ), torch.no_grad(), viewer_utils.SetTrace(self.check_interrupt): - outputs = self.viewer.get_model().get_outputs(camera) + outputs = self.viewer.get_model().get_outputs_for_camera_ray_bundle(camera_ray_bundle,camera=camera) else: with torch.no_grad(), viewer_utils.SetTrace(self.check_interrupt): - outputs = self.viewer.get_model().get_outputs(camera) + outputs = self.viewer.get_model().get_outputs_for_camera_ray_bundle(camera_ray_bundle,camera=camera) except viewer_utils.IOChangeException: self.viewer.get_model().train() raise