Add accumulation

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):

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

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)

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