diff --git a/mattport/nerf/graph/vanilla_nerf.py b/mattport/nerf/graph/vanilla_nerf.py index bd93f9158..52994e88f 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 RGBRenderer +from mattport.nerf.renderers import AccumulationRenderer, RGBRenderer from mattport.nerf.sampler import PDFSampler, UniformSampler from mattport.structures import colors from mattport.structures.rays import RaySamples @@ -112,6 +112,7 @@ class NeRFGraph(Graph): # renderers self.renderer_rgb = RGBRenderer(background_color=colors.WHITE) + self.renderer_accumulation = AccumulationRenderer() # losses self.rgb_loss = MSELoss() @@ -144,6 +145,8 @@ class NeRFGraph(Graph): rgb=coarse_field_outputs[FieldHeadNames.RGB], weights=coarse_weights, ) # RendererOutputs + coarse_renderer_accumulation = self.renderer_accumulation(coarse_weights) # RendererOutputs + # fine network: pdf_ray_samples = self.sampler_pdf(uniform_ray_samples, coarse_weights) # RaySamples fine_field_outputs = self.field_fine(pdf_ray_samples) # FieldOutputs @@ -154,8 +157,15 @@ class NeRFGraph(Graph): rgb=fine_field_outputs[FieldHeadNames.RGB], weights=fine_weights, ) # RendererOutputs + fine_renderer_accumulation = self.renderer_accumulation(fine_weights) # RendererOutputs + # outputs: - outputs = {"rgb_coarse": coarse_renderer_outputs.rgb, "rgb_fine": fine_renderer_outputs.rgb} + outputs = { + "rgb_coarse": coarse_renderer_outputs.rgb, + "rgb_fine": fine_renderer_outputs.rgb, + "accumulation_coarse": coarse_renderer_accumulation.accumulation, + "accumulation_fine": fine_renderer_accumulation.accumulation, + } return outputs def get_losses(self, batch, graph_outputs): diff --git a/mattport/nerf/renderers.py b/mattport/nerf/renderers.py index 929014ce0..8fab1417d 100644 --- a/mattport/nerf/renderers.py +++ b/mattport/nerf/renderers.py @@ -15,6 +15,7 @@ class RendererOutputs: rgb: TensorType["num_rays", 3] = None density: TensorType["num_rays", 1] = None + accumulation: TensorType["num_rays", 1] = None class RGBRenderer(nn.Module): @@ -40,7 +41,7 @@ class RGBRenderer(nn.Module): weights (TensorType[..., "num_samples"]): Weights for each sample Returns: - TensorType[..., "out_dim"]: Composited RGB ray + RendererOutputs: Outputs with rgb values. """ rgb = torch.sum(weights[..., None] * rgb, dim=-2) @@ -55,3 +56,26 @@ class RGBRenderer(nn.Module): renderer_outputs = RendererOutputs(rgb=rgb) return renderer_outputs + + +class AccumulationRenderer(nn.Module): + """Accumulated value along a ray.""" + + @classmethod + def forward( + cls, + weights: TensorType[..., "num_samples"], + ) -> RendererOutputs: + """Composite samples along ray and calculate accumulation. + + Args: + weights (TensorType[..., "num_samples"]): Weights for each sample + + Returns: + RendererOutputs: Outputs with accumulated values. + """ + + accumulation = torch.sum(weights, dim=-1)[..., None] + + renderer_outputs = RendererOutputs(accumulation=accumulation) + return renderer_outputs diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py index 5289c796e..37eaa260e 100644 --- a/mattport/nerf/trainer.py +++ b/mattport/nerf/trainer.py @@ -243,19 +243,28 @@ class Trainer: chunk_size = 1024 rgb_coarse = [] rgb_fine = [] + accumulation_coarse = [] + accumulation_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) rgb_coarse.append(graph_outputs["rgb_coarse"]) rgb_fine.append(graph_outputs["rgb_fine"]) + accumulation_coarse.append(graph_outputs["accumulation_coarse"]) + accumulation_fine.append(graph_outputs["accumulation_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() 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") + self.writer.write_image(f"image_idx_{image_idx}-rgb_coarse_fine", combined_image, step, group="val_img") + + 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") coarse_psnr = get_psnr(image, rgb_coarse) - self.writer.write_scalar(f"image_idx_{image_idx}-coarse_psnr", float(coarse_psnr), step, group="val") + self.writer.write_scalar(f"image_idx_{image_idx}", float(coarse_psnr), step, group="val") fine_psnr = get_psnr(image, rgb_fine) self.stats.update_value(Stats.CURR_TEST_PSNR, float(fine_psnr), step) diff --git a/tests/test_run_train.py b/tests/test_run_train.py index 93a8d99ca..879f9ecdb 100644 --- a/tests/test_run_train.py +++ b/tests/test_run_train.py @@ -2,7 +2,7 @@ Default test to make sure train runs """ import pytest -from hydra.experimental import compose, initialize +from hydra import compose, initialize from scripts.run_train import main