add depth composite, really slow streaming though

nerfstudio/configs/method_configs.py

@@ -71,6 +71,7 @@ from nerfstudio.models.semantic_nerfw import SemanticNerfWModelConfig
 from nerfstudio.models.tensorf import TensoRFModelConfig
 from nerfstudio.models.vanilla_nerf import NeRFModel, VanillaModelConfig
 from nerfstudio.pipelines.base_pipeline import VanillaPipelineConfig
+from nerfstudio.data.datamanagers.full_images_datamanager import FullImageDatamanagerConfig
 from nerfstudio.pipelines.dynamic_batch import DynamicBatchPipelineConfig
 from nerfstudio.plugins.registry import discover_methods
 

nerfstudio/models/gaussian_splatting.py

@@ -611,19 +611,19 @@ class GaussianSplattingModel(Model):
             W,
             background,
         )
-        # depth_im = RasterizeGaussians.apply(
-        #     self.xys[rend_mask],
-        #     depths[rend_mask],
-        #     self.radii[rend_mask],
-        #     conics[rend_mask],
-        #     num_tiles_hit[rend_mask],
-        #     depths[rend_mask,None],
-        #     torch.sigmoid(opacities_crop[rend_mask]),
-        #     H,
-        #     W,
-        #     torch.ones(1,device=self.device)*10,
-        # )[...,0:1]
-        return {"rgb": rgb}#,'depth':depth_im}
+        depth_im = RasterizeGaussians.apply(
+            self.xys[rend_mask],
+            depths[rend_mask],
+            self.radii[rend_mask],
+            conics[rend_mask],
+            num_tiles_hit[rend_mask],
+            depths[rend_mask,None].repeat(1,3),
+            torch.sigmoid(opacities_crop[rend_mask]),
+            H,
+            W,
+            torch.ones(3,device=self.device)*10,
+        )[...,0:1]
+        return {"rgb": rgb,'depth':depth_im}
 
     def get_metrics_dict(self, outputs, batch) -> Dict[str, torch.Tensor]:
         """Compute and returns metrics.

nerfstudio/pipelines/base_pipeline.py

@@ -367,13 +367,7 @@ class VanillaPipeline(Pipeline):
         """
         self.eval()
         metrics_dict_list = []
-<<<<<<< HEAD
-        # TODO (jake-austin): Can we remove this assert statement? Are all the needed functions
-        # abstract functions in the upper-most DataManager class?
-        assert isinstance(self.datamanager, VanillaDataManager)
-=======
         assert isinstance(self.datamanager, (VanillaDataManager, ParallelDataManager))
->>>>>>> 8e2a0651d0dc9bb9d6de5d4a934d0a195ad5032d
         num_images = len(self.datamanager.fixed_indices_eval_dataloader)
         with Progress(
             TextColumn("[progress.description]{task.description}"),

nerfstudio/viewer_beta/render_state_machine.py

@@ -128,9 +128,9 @@ class RenderStateMachine(threading.Thread):
             with self.viewer.train_lock if self.viewer.train_lock is not None else contextlib.nullcontext():
                 if isinstance(self.viewer.get_model(),GaussianSplattingModel):
                     camera_ray_bundle=None
+                    self.viewer.get_model().set_crop(obb)
                 else:
                     camera_ray_bundle = camera.generate_rays(camera_indices=0, obb_box=obb)
-                self.viewer.get_model().set_crop(obb)
                 self.viewer.get_model().eval()
                 step = self.viewer.step
                 try:
@@ -154,14 +154,19 @@ class RenderStateMachine(threading.Thread):
                     self.viewer.get_model().train()
                     raise
                 self.viewer.get_model().train()
-            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
+            if isinstance(self.viewer.get_model(),GaussianSplattingModel):
+                num_rays = image_height*image_width
+            else:
+                num_rays = len(camera_ray_bundle)
+            if self.viewer.control_panel.layer_depth:
+                if isinstance(self.viewer.get_model(),GaussianSplattingModel):
+                    outputs["gl_z_buf_depth"] = outputs['depth']
+                else:
+                    # 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(