more bugfix, number of gaussians doesnt go crazy anymore

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,

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

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."""

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.stop_split_at:
                     #then we densify
-                    avg_grad_norm = (self.xys_grad_norm / self.vis_counts) * max(self.last_size[0],self.last_size[1])
+                    avg_grad_norm = (self.xys_grad_norm / self.vis_counts) * 0.5* max(self.last_size[0],self.last_size[1])
                     high_grads = (avg_grad_norm > 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]:

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)

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()