diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index d59db18c7..7f07900c8 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -62,6 +62,7 @@ from nerfstudio.models.depth_nerfacto import DepthNerfactoModelConfig from nerfstudio.models.gaussian_splatting_inhouse import GaussianSplattingModelConfig from nerfstudio.models.generfacto import GenerfactoModelConfig from nerfstudio.models.instant_ngp import InstantNGPModelConfig +from nerfstudio.data.dataparsers.colmap_dataparser import ColmapDataParserConfig from nerfstudio.models.mipnerf import MipNerfModel from nerfstudio.models.nerfacto import NerfactoModelConfig from nerfstudio.models.neus import NeuSModelConfig @@ -637,13 +638,13 @@ method_configs["gaussian-splatting"] = TrainerConfig( method_name="gaussian-splatting", steps_per_eval_image=10, steps_per_eval_batch=10, - steps_per_save=2000, + steps_per_save=2000000,#saving is broken steps_per_eval_all_images=1000000, # set to a very large model so we don't eval with all images - max_num_iterations=20001, + max_num_iterations=30000, mixed_precision=False, pipeline=VanillaPipelineConfig( datamanager=FullImageDatamanagerConfig( - dataparser=NerfstudioDataParserConfig(), + dataparser=ColmapDataParserConfig(load_3D_points=True,downscale_factor=2), ), model=GaussianSplattingModelConfig(), ), @@ -655,6 +656,13 @@ method_configs["gaussian-splatting"] = TrainerConfig( max_steps=30000, ), }, + "color": { + "optimizer": AdamOptimizerConfig(lr=0.00016, eps=1e-15), + "scheduler": ExponentialDecaySchedulerConfig( + lr_final=0.0000016, + max_steps=30000, + ), + }, "opacity": { "optimizer": AdamOptimizerConfig(lr=0.05, eps=1e-15), "scheduler": None, diff --git a/nerfstudio/data/datamanagers/full_images_datamanager.py b/nerfstudio/data/datamanagers/full_images_datamanager.py index cdb437bd9..107074e2b 100644 --- a/nerfstudio/data/datamanagers/full_images_datamanager.py +++ b/nerfstudio/data/datamanagers/full_images_datamanager.py @@ -289,7 +289,7 @@ class FullImageDatamanager(DataManager, Generic[TDataset]): data["image"] = data["image"].to(self.device) assert len(self.train_dataset.cameras.shape) == 1, "Assumes single batch dimension" - camera = self.train_dataset.cameras[image_idx].to(self.device) + camera = self.train_dataset.cameras[image_idx:image_idx+1].to(self.device) return camera, data def next_eval(self, step: int) -> Tuple[Cameras, Dict]: diff --git a/nerfstudio/models/gaussian_splatting_inhouse.py b/nerfstudio/models/gaussian_splatting_inhouse.py index 259b7bb34..0011a5fd5 100644 --- a/nerfstudio/models/gaussian_splatting_inhouse.py +++ b/nerfstudio/models/gaussian_splatting_inhouse.py @@ -30,8 +30,15 @@ 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 + +from diff_rast.rasterize import RasterizeGaussians +from diff_rast.project_gaussians import ProjectGaussians + + def random_quat_tensor(N, **kwargs): u = torch.rand(N, **kwargs) v = torch.rand(N, **kwargs) @@ -54,9 +61,9 @@ def identity_quat(N, **kwargs): def projection_matrix(znear, zfar, fovx, fovy, **kwargs): - t = math.tan(0.5 * fovy) + t = znear * math.tan(0.5 * fovy) b = -t - r = math.tan(0.5 * fovx) + r = znear * math.tan(0.5 * fovx) l = -r n = znear f = zfar @@ -64,8 +71,8 @@ def projection_matrix(znear, zfar, fovx, fovy, **kwargs): [ [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), -2.0 * f * n / (f - n)], - [0.0, 0.0, -1.0, 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 ) @@ -75,7 +82,14 @@ def projection_matrix(znear, zfar, fovx, fovy, **kwargs): class GaussianSplattingModelConfig(ModelConfig): """Gaussian Splatting Model Config""" _target: Type = field(default_factory=lambda: GaussianSplattingModel) - num_points: int = 100000 + warmup_length:int = 300 + """period of steps where resolution is upscaled iteratively""" + refine_every:int = 300 + """period of steps where gaussians are culled and densified""" + cull_threshold:float = .005 + """threshold of opacity for culling gaussians""" + reset_alpha_every:int = 10 + """Every this many refinement steps, reset the alpha""" class GaussianSplattingModel(Model): @@ -88,26 +102,136 @@ class GaussianSplattingModel(Model): """ config: GaussianSplattingModelConfig + def __init__(self,*args,**kwargs): + self.seed_pts = kwargs['seed_points'] + 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 = 2 * (torch.nn.Parameter(torch.rand(self.config.num_points, 3, device=self.device) - 0.5)) - self.scales = torch.nn.Parameter(torch.rand(self.config.num_points, 3, device=self.device)) - # self.quats = random_quat_tensor(self.num_points, device=self.device) - self.quats = torch.nn.Parameter(identity_quat(self.num_points, device=self.device)) - self.rgbs = torch.nn.Parameter(torch.rand(self.config.num_points, 3, device=self.device)) - self.opacities = torch.nn.Parameter(0.5 * torch.ones(self.config.num_points, 1, device=self.device)) - + self.means = torch.nn.Parameter(self.seed_pts[0]) + init_scale = torch.log(torch.tensor(.002)).item() + self.scales = torch.nn.Parameter(torch.full((self.num_points,3),init_scale)) + self.quats = torch.nn.Parameter(random_quat_tensor(self.num_points)) + rgbinit = torch.clamp((self.seed_pts[1].float())/255.,0,1) + self.rgbs = torch.nn.Parameter(torch.logit(rgbinit)) + self.opacities = torch.nn.Parameter(torch.zeros(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 + + 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] + param_state['exp_avg'] = param_state['exp_avg'][~deleted_mask.squeeze()] + param_state['exp_avg_sq'] = param_state['exp_avg_sq'][~deleted_mask.squeeze()] + del optimizer.state[param] + optimizer.state[param] = param_state + optimizer.param_groups[0]['params'] = new_params + def duplicate_ids(self,ids): + """Duplicates the gaussians matching ids with a small random delta""" + with torch.no_grad(): + newmeans = self.means[ids] + torch.randn_like(self.means[ids])*.001 + newscales = self.scales[ids].clone() + newquats = self.quats[ids].clone() + newrgbs = self.rgbs[ids].clone() + newopacities = self.opacities[ids].clone() + self.means.data = torch.cat([self.means.data,newmeans.data],dim=0) + self.scales.data = torch.cat([self.scales.data,newscales.data],dim=0) + self.quats.data = torch.cat([self.quats.data,newquats.data],dim=0) + self.rgbs.data = torch.cat([self.rgbs.data,newrgbs.data],dim=0) + self.opacities.data = torch.cat([self.opacities.data,newopacities.data],dim=0) + + def refinement_before(self, optimizers:Optimizers, step): + print("Inside refinement before") + if self.step > self.config.warmup_length: + # 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: + #then we densify + self.split_gaussians() + self.dup_gaussians() + if self.step // self.config.refine_every % self.config.reset_alpha_every == 0: + print("Resetting alpha") + with torch.no_grad(): + reset_value = .01 + self.opacities.data = torch.full_like(self.opacities.data,torch.logit(torch.tensor(reset_value)).item()) + + def cull_gaussians(self): + """ + This function deletes gaussians with under a certain opacity threshold + """ + n_bef = self.num_points + with torch.no_grad(): + culls = (torch.sigmoid(self.opacities) < self.config.cull_threshold).squeeze() + self.means = Parameter(self.means[~culls]) + self.scales = Parameter(self.scales[~culls]) + self.quats = Parameter(self.quats[~culls]) + self.rgbs = Parameter(self.rgbs[~culls]) + self.opacities = Parameter(self.opacities[~culls]) + print(f"Culled {n_bef - self.num_points} gaussians") + return culls + + def split_gaussians(self): + """ + This function splits gaussians that are too large + """ + pass + + def dup_gaussians(self): + """ + This function duplicates gaussians that are too small + """ + pass + + + @property + def num_points(self): + return self.means.shape[0] + def get_training_callbacks( self, training_callback_attributes: TrainingCallbackAttributes ) -> List[TrainingCallback]: + """ + List of things that are important: + every 300 steps, densify by calculating positional gradients and duplicating/splitting + every 3000 steps: drop all alpha to low, and continue optimization + cull gaussians with alpha under a threshold + upscale image resolution every 300 iterations + SH coefficients are masked initially and get activated over time - return [] + + """ + 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] + )) + 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 @@ -123,7 +247,7 @@ class GaussianSplattingModel(Model): "rotation": [self.quats], } - def get_outputs(self, ray_bundle: Cameras) -> Dict[str, Union[torch.Tensor, List]]: + def get_outputs(self, camera: Cameras) -> Dict[str, Union[torch.Tensor, List]]: """Takes in a Ray Bundle and returns a dictionary of outputs. TODO (jake-austin): use the new homebrew nerfstudio gaussian rasterization code instead @@ -135,16 +259,63 @@ class GaussianSplattingModel(Model): 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" - #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() - # return {"rgb": rgb} + #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) + #vickie pops off here + 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 + fovx = 2 * math.atan(camera.width / (2 * camera.fx)) + fovy = 2 * math.atan(camera.height / (2 * camera.fy)) + projmat = projection_matrix(.0001,1000,fovx,fovy).to(self.device) + W, H = camera.width.item(), camera.height.item() + BLOCK_X, BLOCK_Y = 16, 16 + tile_bounds = ( + (W + BLOCK_X - 1) // BLOCK_X, + (H + BLOCK_Y - 1) // BLOCK_Y, + 1, + ) + background = torch.ones(3, device=self.device) + xys, depths, radii, conics, num_tiles_hit = ProjectGaussians.apply( + self.means, + torch.exp(self.scales), + 1, + self.quats, + viewmat.squeeze()[:3,:], + projmat.squeeze()@viewmat.squeeze(), + camera.fx.item(), + camera.fy.item(), + H, + W, + tile_bounds + ) + rgb = RasterizeGaussians.apply( + xys, + depths, + radii, + conics, + num_tiles_hit, + torch.sigmoid(self.rgbs), + torch.sigmoid(self.opacities), + H, + W, + background, + ) + return {"rgb": rgb} def get_metrics_dict(self, outputs, batch) -> Dict[str, torch.Tensor]: @@ -155,7 +326,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. @@ -165,8 +336,9 @@ 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)) - Ll1 = torch.nn.functional.l1_loss(batch['image'], outputs['rgb']) + Ll1 = torch.nn.functional.mse_loss(batch['image'], outputs['rgb']) + # This simloss makes the results look weird, removing for now + # simloss = self.ssim(batch['image'].permute(2,0,1)[None,...], outputs['rgb'].permute(2,0,1)[None,...]) return {"main_loss": Ll1} @torch.no_grad() diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index 6cd03fee9..1ee3bb49d 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/nerfstudio/pipelines/base_pipeline.py @@ -260,7 +260,8 @@ class VanillaPipeline(Pipeline): device=device, test_mode=test_mode, world_size=world_size, local_rank=local_rank ) print("Inside base pipeline") - import pdb;pdb.set_trace() + pts = self.datamanager.train_dataparser_outputs.metadata['points3D_xyz'] + pts_rgb = self.datamanager.train_dataparser_outputs.metadata['points3D_rgb'] self.datamanager.to(device) # TODO(ethan): get rid of scene_bounds from the model assert self.datamanager.train_dataset is not None, "Missing input dataset" @@ -271,6 +272,7 @@ class VanillaPipeline(Pipeline): metadata=self.datamanager.train_dataset.metadata, device=device, grad_scaler=grad_scaler, + seed_points = (pts,pts_rgb) ) self.model.to(device) diff --git a/nerfstudio/viewer_beta/render_state_machine.py b/nerfstudio/viewer_beta/render_state_machine.py index d6aedf7de..07c32ea77 100644 --- a/nerfstudio/viewer_beta/render_state_machine.py +++ b/nerfstudio/viewer_beta/render_state_machine.py @@ -124,20 +124,9 @@ class RenderStateMachine(threading.Thread): camera = camera.to(self.viewer.get_model().device) assert camera is not None, "render called before viewer connected" -<<<<<<< HEAD - with self.viewer.train_lock if self.viewer.train_lock is not None else contextlib.nullcontext(): - #TODO jake-austin: Make this check whether the model inherits from a camera based model or a ray based model - if True: - pass - else: - camera_ray_bundle = camera.generate_rays(camera_indices=0, aabb_box=self.viewer.get_model().render_aabb) - - with TimeWriter(None, None, write=False) as vis_t: -======= with TimeWriter(None, None, write=False) as vis_t: with self.viewer.train_lock if self.viewer.train_lock is not None else contextlib.nullcontext(): camera_ray_bundle = camera.generate_rays(camera_indices=0, obb_box=obb) ->>>>>>> 8812d7a5b092beb822ddd2b191c5e82cd3f66ed5 self.viewer.get_model().eval() step = self.viewer.step try: @@ -153,44 +142,22 @@ 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_for_camera_ray_bundle(camera_ray_bundle) + outputs = self.viewer.get_model().get_outputs(camera) else: -<<<<<<< HEAD - background_color = torch.tensor( - [color[0] / 255.0, color[1] / 255.0, color[2] / 255.0], - device=self.viewer.get_model().device, - ) - with background_color_override_context(background_color), torch.no_grad(): - if True: - outputs = self.viewer.get_model().get_outputs_for_camera_ray_bundle(camera) - else: - outputs = self.viewer.get_model().get_outputs_for_camera_ray_bundle(camera_ray_bundle) - else: - with torch.no_grad(): - if True: - outputs = self.viewer.get_model().get_outputs_for_camera_ray_bundle(camera) - else: - outputs = self.viewer.get_model().get_outputs_for_camera_ray_bundle(camera_ray_bundle) - self.viewer.get_model().train() - if True: - num_rays = (camera.height * camera.width).item() - else: - num_rays = len(camera_ray_bundle) -======= with torch.no_grad(), viewer_utils.SetTrace(self.check_interrupt): - outputs = self.viewer.get_model().get_outputs_for_camera_ray_bundle(camera_ray_bundle) + outputs = self.viewer.get_model().get_outputs(camera) except viewer_utils.IOChangeException: self.viewer.get_model().train() raise self.viewer.get_model().train() - num_rays = len(camera_ray_bundle) - if self.viewer.control_panel.layer_depth: - # convert to z_depth if depth compositing is enabled - R = camera.camera_to_worlds[0:3, 0:3].T - pts = camera_ray_bundle.directions * outputs["depth"] - pts = (R @ (pts.view(-1, 3).T)).T.view(*camera_ray_bundle.directions.shape) - outputs["gl_z_buf_depth"] = -pts[..., 2:3] # negative z axis is the coordinate convention ->>>>>>> 8812d7a5b092beb822ddd2b191c5e82cd3f66ed5 + num_rays = image_height*image_width + # num_rays = len(camera_ray_bundle) + # if self.viewer.control_panel.layer_depth: + # # convert to z_depth if depth compositing is enabled + # R = camera.camera_to_worlds[0:3, 0:3].T + # pts = camera_ray_bundle.directions * outputs["depth"] + # pts = (R @ (pts.view(-1, 3).T)).T.view(*camera_ray_bundle.directions.shape) + # outputs["gl_z_buf_depth"] = -pts[..., 2:3] # negative z axis is the coordinate convention render_time = vis_t.duration if writer.is_initialized(): writer.put_time( diff --git a/nerfstudio/viewer_beta/utils.py b/nerfstudio/viewer_beta/utils.py index a3dd7b1c2..429edcb05 100644 --- a/nerfstudio/viewer_beta/utils.py +++ b/nerfstudio/viewer_beta/utils.py @@ -65,7 +65,7 @@ def get_camera( cx=pp_w, cy=pp_h, camera_type=camera_type, - camera_to_worlds=camera_state.c2w.to(torch.float32), + camera_to_worlds=camera_state.c2w.to(torch.float32)[None,...], times=torch.tensor([0.0], dtype=torch.float32), ) return camera diff --git a/nerfstudio/viewer_beta/viewer.py b/nerfstudio/viewer_beta/viewer.py index 6e6e02a29..cb97be63a 100644 --- a/nerfstudio/viewer_beta/viewer.py +++ b/nerfstudio/viewer_beta/viewer.py @@ -211,6 +211,7 @@ class Viewer: def update_camera_poses(self): # Update the train camera locations based on optimization + return assert self.camera_handles is not None idxs = list(self.camera_handles.keys()) if hasattr(self.pipeline.datamanager, "train_camera_optimizer"):