From b350a68ff4d6289074220044b1359b855eb722c5 Mon Sep 17 00:00:00 2001 From: Justin Kerr Date: Sat, 30 Sep 2023 11:01:45 -0700 Subject: [PATCH] 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