From 6e20b7b41899f250607c42419b784ae349e4010a Mon Sep 17 00:00:00 2001 From: Justin Kerr Date: Thu, 12 Oct 2023 18:36:59 -0700 Subject: [PATCH] add depth composite, really slow streaming though --- nerfstudio/configs/method_configs.py | 1 + nerfstudio/models/gaussian_splatting.py | 26 +++++++++---------- nerfstudio/pipelines/base_pipeline.py | 6 ----- .../viewer_beta/render_state_machine.py | 23 +++++++++------- 4 files changed, 28 insertions(+), 28 deletions(-) diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index dcfe5599c..72469180e 100644 --- a/nerfstudio/configs/method_configs.py +++ b/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 diff --git a/nerfstudio/models/gaussian_splatting.py b/nerfstudio/models/gaussian_splatting.py index 00908e720..35d416b33 100644 --- a/nerfstudio/models/gaussian_splatting.py +++ b/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. diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index 85cc7e528..b3f27b91d 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/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}"), diff --git a/nerfstudio/viewer_beta/render_state_machine.py b/nerfstudio/viewer_beta/render_state_machine.py index 79e1cd9ca..9283b0e0a 100644 --- a/nerfstudio/viewer_beta/render_state_machine.py +++ b/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(