From 17e95e1ba2a812b29a56e20469a2ff26b3214a1b Mon Sep 17 00:00:00 2001 From: Justin Kerr Date: Tue, 10 Oct 2023 10:29:44 -0700 Subject: [PATCH] more bugfix, number of gaussians doesnt go crazy anymore --- nerfstudio/configs/method_configs.py | 2 +- .../datamanagers/full_images_datamanager.py | 13 ++--- .../data/dataparsers/colmap_dataparser.py | 2 +- .../models/gaussian_splatting_inhouse.py | 56 +++++++++---------- nerfstudio/pipelines/base_pipeline.py | 2 +- nerfstudio/viewer_beta/viewer.py | 7 ++- 6 files changed, 39 insertions(+), 43 deletions(-) diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 1bf9324b5..33557ece2 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -603,7 +603,7 @@ method_configs["gaussian-splatting"] = TrainerConfig( ), optimizers={ "xyz": { - "optimizer": AdamOptimizerConfig(lr=1.6e-4, eps=1e-15), + "optimizer": AdamOptimizerConfig(lr=1e-4, eps=1e-15), "scheduler": ExponentialDecaySchedulerConfig( lr_final=1.6e-6, max_steps=15000, diff --git a/nerfstudio/data/datamanagers/full_images_datamanager.py b/nerfstudio/data/datamanagers/full_images_datamanager.py index 107074e2b..0a200b469 100644 --- a/nerfstudio/data/datamanagers/full_images_datamanager.py +++ b/nerfstudio/data/datamanagers/full_images_datamanager.py @@ -46,8 +46,8 @@ from nerfstudio.configs.dataparser_configs import AnnotatedDataParserUnion from nerfstudio.data.datamanagers.base_datamanager import (DataManager, DataManagerConfig,TDataset) from nerfstudio.data.dataparsers.base_dataparser import DataparserOutputs -from nerfstudio.data.dataparsers.blender_dataparser import \ - BlenderDataParserConfig +from nerfstudio.data.dataparsers.nerfstudio_dataparser import \ + NerfstudioDataParserConfig from nerfstudio.data.datasets.base_dataset import InputDataset from nerfstudio.utils.misc import get_orig_class from nerfstudio.utils.rich_utils import CONSOLE @@ -57,7 +57,7 @@ CONSOLE = Console(width=120) @dataclass class FullImageDatamanagerConfig(DataManagerConfig): _target: Type = field(default_factory=lambda: FullImageDatamanager) - dataparser: AnnotatedDataParserUnion = BlenderDataParserConfig() + dataparser: AnnotatedDataParserUnion = NerfstudioDataParserConfig() camera_res_scale_factor: float = 1.0 """The scale factor for scaling spatial data such as images, mask, semantics along with relevant information about camera intrinsics @@ -75,11 +75,8 @@ class FullImageDatamanagerConfig(DataManagerConfig): class FullImageDatamanager(DataManager, Generic[TDataset]): """ A datamanager that outputs full images and cameras instead of raybundles. This makes the - datamanager far more lightweight since we dont have to do ray generation for things like - gaussian splatting that dont require it. - - TODO (jake-austin): Figure out why the dataparser doesnt show up as a subcommand, preventing us - from specifying nerfstudio-data instead of blender-data + datamanager more lightweight since we don't have to do generate rays. Useful for full-image + training e.g. rasterization pipelines """ config: FullImageDatamanagerConfig diff --git a/nerfstudio/data/dataparsers/colmap_dataparser.py b/nerfstudio/data/dataparsers/colmap_dataparser.py index c68a9cf6a..27ff0c55c 100644 --- a/nerfstudio/data/dataparsers/colmap_dataparser.py +++ b/nerfstudio/data/dataparsers/colmap_dataparser.py @@ -70,7 +70,7 @@ class ColmapDataParserConfig(DataParserConfig): """Path to masks directory. If not set, masks are not loaded.""" depths_path: Optional[Path] = None """Path to depth maps directory. If not set, depths are not loaded.""" - colmap_path: Path = Path("sparse/0") + colmap_path: Path = Path("colmap/sparse/0") """Path to the colmap reconstruction directory relative to the data path.""" load_3D_points: bool = False """Whether to load the 3D points from the colmap reconstruction.""" diff --git a/nerfstudio/models/gaussian_splatting_inhouse.py b/nerfstudio/models/gaussian_splatting_inhouse.py index f627b92fb..aa0d15a92 100644 --- a/nerfstudio/models/gaussian_splatting_inhouse.py +++ b/nerfstudio/models/gaussian_splatting_inhouse.py @@ -100,7 +100,7 @@ class GaussianSplattingModelConfig(ModelConfig): """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 + max_screen_size: int = 100 """maximum screen size of a gaussian, in pixels""" random_init: bool = False """whether to initialize the positions uniformly randomly (not SFM points)""" @@ -108,6 +108,8 @@ class GaussianSplattingModelConfig(ModelConfig): """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 @@ -142,8 +144,7 @@ class GaussianSplattingModel(Model): self.scales = torch.nn.Parameter(torch.log(avg_dist.repeat(1, 3))) self.quats = torch.nn.Parameter(random_quat_tensor(self.num_points)) - self.degree = 3 - dim_sh = num_sh_bases(self.degree) + 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) @@ -181,7 +182,7 @@ class GaussianSplattingModel(Model): 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.degree)-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) @@ -261,7 +262,7 @@ class GaussianSplattingModel(Model): 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 @@ -284,7 +285,6 @@ class GaussianSplattingModel(Model): 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] - # add_idcs = torch.cat([split_idcs, dup_idcs], dim=0) param_groups = self.get_param_groups() for group,param in param_groups.items(): @@ -318,10 +318,10 @@ class GaussianSplattingModel(Model): n_bef = self.num_points # cull transparent ones culls = (torch.sigmoid(self.opacities) < self.config.cull_alpha_thresh).squeeze() - # cull huge ones - culls = culls | (torch.exp(self.scales).max(dim=-1).values > self.config.cull_scale_thresh).squeeze() - # cull big screen space 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()) @@ -342,7 +342,8 @@ class GaussianSplattingModel(Model): 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 - rots = quat_to_rotmat(self.quats[split_mask].repeat(samps,1)) # 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 @@ -352,8 +353,8 @@ class GaussianSplattingModel(Model): # step 3, sample new opacities new_opacities = self.opacities[split_mask].repeat(samps,1) # step 4, sample new scales - new_scales = torch.logit(torch.exp(self.scales[split_mask])/1.6).repeat(samps,1) - self.scales[split_mask] = torch.logit(torch.exp(self.scales[split_mask])/1.6) + 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 @@ -496,10 +497,10 @@ class GaussianSplattingModel(Model): #Important to allow xys grads to populate properly if self.training: self.xys.retain_grad() - if self.degree > 0: + 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.degree) + 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: @@ -518,21 +519,18 @@ class GaussianSplattingModel(Model): W, background, ) - if not self.training: - 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] - else: - depth_im = None + 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]: diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index 1ee3bb49d..1b802380a 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/nerfstudio/pipelines/base_pipeline.py @@ -259,7 +259,7 @@ class VanillaPipeline(Pipeline): self.datamanager: DataManager = config.datamanager.setup( device=device, test_mode=test_mode, world_size=world_size, local_rank=local_rank ) - print("Inside base pipeline") + #TODO make this work with normal nerfacto too pts = self.datamanager.train_dataparser_outputs.metadata['points3D_xyz'] pts_rgb = self.datamanager.train_dataparser_outputs.metadata['points3D_rgb'] self.datamanager.to(device) diff --git a/nerfstudio/viewer_beta/viewer.py b/nerfstudio/viewer_beta/viewer.py index cb97be63a..b76c2dd53 100644 --- a/nerfstudio/viewer_beta/viewer.py +++ b/nerfstudio/viewer_beta/viewer.py @@ -211,13 +211,14 @@ 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"): camera_optimizer = self.pipeline.datamanager.train_camera_optimizer - else: + elif hasattr(self.pipeline.model, "camera_optimizer"): camera_optimizer = self.pipeline.model.camera_optimizer + else: + return + idxs = list(self.camera_handles.keys()) with torch.no_grad(): assert isinstance(camera_optimizer, CameraOptimizer) c2ws_delta = camera_optimizer(torch.tensor(idxs, device=camera_optimizer.device)).cpu().numpy()