diff --git a/mattport/nerf/graph/vanilla_nerf.py b/mattport/nerf/graph/vanilla_nerf.py index 7c0bf5600..13db5b7d7 100644 --- a/mattport/nerf/graph/vanilla_nerf.py +++ b/mattport/nerf/graph/vanilla_nerf.py @@ -16,7 +16,7 @@ from mattport.nerf.field_modules.mlp import MLP from mattport.nerf.field_modules.ray_generator import RayGenerator from mattport.nerf.graph.base import Graph from mattport.nerf.loss import MSELoss -from mattport.nerf.renderers import AccumulationRenderer, DisparityRenderer, RGBRenderer +from mattport.nerf.renderers import AccumulationRenderer, DepthRenderer, RGBRenderer from mattport.nerf.sampler import PDFSampler, UniformSampler from mattport.structures import colors from mattport.structures.rays import RaySamples @@ -114,7 +114,7 @@ class NeRFGraph(Graph): # renderers self.renderer_rgb = RGBRenderer(background_color=colors.WHITE) self.renderer_accumulation = AccumulationRenderer() - self.renderer_disparity = DisparityRenderer() + self.renderer_depth = DepthRenderer() # losses self.rgb_loss = MSELoss() @@ -148,7 +148,7 @@ class NeRFGraph(Graph): weights=coarse_weights, ) # RendererOutputs coarse_renderer_accumulation = self.renderer_accumulation(coarse_weights) # RendererOutputs - coarse_renderer_disparity = self.renderer_disparity(coarse_weights, uniform_ray_samples.ts) + coarse_renderer_depth = self.renderer_depth(coarse_weights, uniform_ray_samples.ts) # fine network: pdf_ray_samples = self.sampler_pdf(uniform_ray_samples, coarse_weights) # RaySamples @@ -161,9 +161,21 @@ class NeRFGraph(Graph): weights=fine_weights, ) # RendererOutputs fine_renderer_accumulation = self.renderer_accumulation(fine_weights) # RendererOutputs - fine_renderer_disparity = self.renderer_disparity(fine_weights, pdf_ray_samples.ts) + fine_renderer_depth = self.renderer_depth(fine_weights, pdf_ray_samples.ts) # TODO refactor this into "vis" section. Doesn't need to be run during training. + coarse_renderer_depth = visualization.apply_depth_colormap( + coarse_renderer_depth.depth, + accumulation=coarse_renderer_accumulation.accumulation, + near_plane=self.near_plane, + far_plane=self.far_plane, + ) + fine_renderer_depth = visualization.apply_depth_colormap( + fine_renderer_depth.depth, + accumulation=fine_renderer_accumulation.accumulation, + near_plane=self.near_plane, + far_plane=self.far_plane, + ) coarse_renderer_accumulation = visualization.apply_colormap(coarse_renderer_accumulation.accumulation) fine_renderer_accumulation = visualization.apply_colormap(fine_renderer_accumulation.accumulation) @@ -173,8 +185,8 @@ class NeRFGraph(Graph): "rgb_fine": fine_renderer_outputs.rgb, "accumulation_coarse": coarse_renderer_accumulation, "accumulation_fine": fine_renderer_accumulation, - "disparity_coarse": coarse_renderer_disparity.disparity, - "disparity_fine": fine_renderer_disparity.disparity, + "depth_coarse": coarse_renderer_depth, + "depth_fine": fine_renderer_depth, } return outputs diff --git a/mattport/nerf/renderers.py b/mattport/nerf/renderers.py index 423e1892d..9175007db 100644 --- a/mattport/nerf/renderers.py +++ b/mattport/nerf/renderers.py @@ -16,7 +16,7 @@ class RendererOutputs: rgb: TensorType["num_rays", 3] = None density: TensorType["num_rays", 1] = None accumulation: TensorType["num_rays", 1] = None - disparity: TensorType["num_rays", 1] = None + depth: TensorType["num_rays", 1] = None class RGBRenderer(nn.Module): @@ -82,7 +82,7 @@ class AccumulationRenderer(nn.Module): return renderer_outputs -class DisparityRenderer(nn.Module): +class DepthRenderer(nn.Module): """Calcualte depth along ray.""" def __init__(self, method: str = "expected") -> None: @@ -103,14 +103,15 @@ class DisparityRenderer(nn.Module): ts (TensorType[..., "num_samples"]): Sample locations along rays Returns: - RendererOutputs: Outputs with disparity values. + RendererOutputs: Outputs with depth values. """ if self.method == "expected": - depth = torch.sum(weights * ts, dim=-1) eps = 1e-10 - disparity = 1.0 / torch.max(eps * torch.ones_like(depth), depth / torch.sum(weights, -1)) + depth = torch.sum(weights * ts, dim=-1) / (torch.sum(weights, -1) + eps) - return RendererOutputs(disparity=disparity) + depth = torch.clip(depth, ts[..., 0], ts[..., -1]) + + return RendererOutputs(depth=depth[..., None]) raise NotImplementedError(f"Method {self.method} not implemented") diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py index e4e9a938c..e816244c0 100644 --- a/mattport/nerf/trainer.py +++ b/mattport/nerf/trainer.py @@ -259,8 +259,8 @@ class Trainer: rgb_fine = [] accumulation_coarse = [] accumulation_fine = [] - disparity_coarse = [] - disparity_fine = [] + depth_coarse = [] + depth_fine = [] for i in range(0, num_rays, chunk_size): ray_indices = all_ray_indices[i : i + chunk_size].to(self.device) graph_outputs = self.graph(ray_indices) @@ -268,14 +268,14 @@ class Trainer: rgb_fine.append(graph_outputs["rgb_fine"]) accumulation_coarse.append(graph_outputs["accumulation_coarse"]) accumulation_fine.append(graph_outputs["accumulation_fine"]) - disparity_coarse.append(graph_outputs["disparity_coarse"]) - disparity_fine.append(graph_outputs["disparity_fine"]) + depth_coarse.append(graph_outputs["depth_coarse"]) + depth_fine.append(graph_outputs["depth_fine"]) rgb_coarse = torch.cat(rgb_coarse).view(image_height, image_width, 3).detach().cpu() rgb_fine = torch.cat(rgb_fine).view(image_height, image_width, 3).detach().cpu() accumulation_coarse = torch.cat(accumulation_coarse).view(image_height, image_width, 3).detach().cpu() accumulation_fine = torch.cat(accumulation_fine).view(image_height, image_width, 3).detach().cpu() - disparity_coarse = torch.cat(disparity_coarse).view(image_height, image_width, 1).detach().cpu() - disparity_fine = torch.cat(disparity_fine).view(image_height, image_width, 1).detach().cpu() + depth_coarse = torch.cat(depth_coarse).view(image_height, image_width, 3).detach().cpu() + depth_fine = torch.cat(depth_fine).view(image_height, image_width, 3).detach().cpu() combined_image = torch.cat([image, rgb_coarse, rgb_fine], dim=1) writer.write_event( @@ -287,10 +287,8 @@ class Trainer: {"name": f"image_idx_{image_idx}", "x": combined_image, "step": step, "group": "val_accumulation"} ) - combined_image = torch.cat([disparity_coarse, disparity_fine], dim=1) - writer.write_event( - {"name": f"image_idx_{image_idx}", "x": combined_image, "step": step, "group": "val_disparity"} - ) + combined_image = torch.cat([depth_coarse, depth_fine], dim=1) + writer.write_event({"name": f"image_idx_{image_idx}", "x": combined_image, "step": step, "group": "val_depth"}) coarse_psnr = get_psnr(image, rgb_coarse) writer.write_event( diff --git a/mattport/utils/writer.py b/mattport/utils/writer.py index bee18aeba..94a6a167d 100644 --- a/mattport/utils/writer.py +++ b/mattport/utils/writer.py @@ -27,7 +27,6 @@ def write_event(args: Dict[str, Any]): assert comms.is_main_process(), "Writing out with process other than main" for writer in EVENT_WRITERS: if "x" in args: - print(args["x"].shape) writer.write_image(**args) elif "scalar" in args: writer.write_scalar(**args)