diff --git a/mattport/nerf/graph/vanilla_nerf.py b/mattport/nerf/graph/vanilla_nerf.py index 52994e88f..16014b239 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, RGBRenderer +from mattport.nerf.renderers import AccumulationRenderer, DisparityRenderer, RGBRenderer from mattport.nerf.sampler import PDFSampler, UniformSampler from mattport.structures import colors from mattport.structures.rays import RaySamples @@ -113,6 +113,7 @@ class NeRFGraph(Graph): # renderers self.renderer_rgb = RGBRenderer(background_color=colors.WHITE) self.renderer_accumulation = AccumulationRenderer() + self.renderer_disparity = DisparityRenderer() # losses self.rgb_loss = MSELoss() @@ -146,6 +147,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) # fine network: pdf_ray_samples = self.sampler_pdf(uniform_ray_samples, coarse_weights) # RaySamples @@ -158,6 +160,7 @@ 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) # outputs: outputs = { @@ -165,6 +168,8 @@ class NeRFGraph(Graph): "rgb_fine": fine_renderer_outputs.rgb, "accumulation_coarse": coarse_renderer_accumulation.accumulation, "accumulation_fine": fine_renderer_accumulation.accumulation, + "disparity_coarse": coarse_renderer_disparity.disparity, + "disparity_fine": fine_renderer_disparity.disparity, } return outputs diff --git a/mattport/nerf/renderers.py b/mattport/nerf/renderers.py index 8fab1417d..423e1892d 100644 --- a/mattport/nerf/renderers.py +++ b/mattport/nerf/renderers.py @@ -16,6 +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 class RGBRenderer(nn.Module): @@ -79,3 +80,37 @@ class AccumulationRenderer(nn.Module): renderer_outputs = RendererOutputs(accumulation=accumulation) return renderer_outputs + + +class DisparityRenderer(nn.Module): + """Calcualte depth along ray.""" + + def __init__(self, method: str = "expected") -> None: + """ + Args: + method (str, optional): Depth calculation method. Defaults to 'expected'. + """ + super().__init__() + if method not in {"expected"}: + raise ValueError(f"{method} is an invalid depth calculation method") + self.method = method + + def forward(self, weights: TensorType[..., "num_samples"], ts: TensorType[..., "num_samples"]) -> RendererOutputs: + """Composite samples along ray and calculate disparities. + + Args: + weights (TensorType[..., "num_samples"]): Weights for each sample + ts (TensorType[..., "num_samples"]): Sample locations along rays + + Returns: + RendererOutputs: Outputs with disparity 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)) + + return RendererOutputs(disparity=disparity) + + raise NotImplementedError(f"Method {self.method} not implemented") diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py index 29a4d0278..5ef5c149c 100644 --- a/mattport/nerf/trainer.py +++ b/mattport/nerf/trainer.py @@ -247,6 +247,8 @@ class Trainer: rgb_fine = [] accumulation_coarse = [] accumulation_fine = [] + disparity_coarse = [] + disparity_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) @@ -254,10 +256,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"]) 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, 1).detach().cpu() accumulation_fine = torch.cat(accumulation_fine).view(image_height, image_width, 1).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() combined_image = torch.cat([image, rgb_coarse, rgb_fine], dim=1) self.writer.write_image(f"image_idx_{image_idx}-rgb_coarse_fine", combined_image, step, group="val_img") @@ -265,6 +271,9 @@ class Trainer: combined_image = torch.cat([accumulation_coarse, accumulation_fine], dim=1) self.writer.write_image(f"image_idx_{image_idx}", combined_image, step, group="val_accumulation") + combined_image = torch.cat([disparity_coarse, disparity_fine], dim=1) + self.writer.write_image(f"image_idx_{image_idx}", combined_image, step, group="val_disparity") + coarse_psnr = get_psnr(image, rgb_coarse) self.writer.write_scalar(f"image_idx_{image_idx}", float(coarse_psnr), step, group="val")