inhouse impl start

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,

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]:

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

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)
 

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(

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

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"):