Merge pull request #20 from ethanweber/tancik/renderers_test

Fix renderer, add test

.github/workflows/code_checks.yml

@@ -31,7 +31,7 @@ jobs:
         options: "--check --verbose --line-length=120"
     - name: Python Pylint
       run: |
-        pylint mattport --max-line-length=120 --generated-members=numpy.*,torch.*,cv2.*,cv.* --good-names=e,c,i,j,x,y,lr,K,fx,fy,cx,cy,H,W --disable=too-many-arguments
+        pylint mattport --max-line-length=120 --generated-members=numpy.*,torch.*,cv2.*,cv.* --good-names=e,c,i,j,x,y,lr,K,fx,fy,cx,cy,H,W --disable=too-many-arguments --disable=fixme
     - name: Test with pytest
       run: |
         pytest tests --ignore tests/nerf/test_ray_generator.py

.vscode/settings.json

@@ -26,6 +26,7 @@
         "--generated-members=numpy.*,torch.*,cv2.*,cv.*",
         "--max-line-length=120",
         "--good-names=e,c,i,j,x,y,lr,K,fx,fy,cx,cy,H,W",
-        "--disable=too-many-arguments"
+        "--disable=too-many-arguments",
+        "--disable=fixme"
     ],
 }

mattport/nerf/renderers.py

@@ -13,7 +13,7 @@ from mattport.nerf.field_modules import field_outputs
 class Renderer(nn.Module):
     """Base Renderer. Intended to be subclassed"""
 
-    def required_render_heads(self) -> Set[field_outputs.FieldOutput]:
+    def required_field_outputs(self) -> Set[field_outputs.FieldOutput]:
         """
         Returns:
             Set[RenderHead]: RenderHeads required for this renderer.
@@ -22,7 +22,7 @@ class Renderer(nn.Module):
 
     def forward(
         self,
-        render_head_outputs: Dict["str", TensorType[..., "num_samples", "features"]],
+        field_output_dict: Dict["str", TensorType[..., "num_samples", "features"]],
         deltas: TensorType[..., "num_samples"],
     ) -> TensorType[..., "out_dim"]:
         """Composite samples along ray and render image
@@ -45,12 +45,12 @@ class RGB(Renderer):
         self.rgb_name = rgb_name
         self.density_name = density_name
 
-    def required_render_heads(self) -> Set[field_outputs.FieldOutput]:
+    def required_field_outputs(self) -> Set[field_outputs.FieldOutput]:
         return set(field_outputs.RGBFieldOutput, field_outputs.DensityFieldOutput)
 
     def forward(
         self,
-        render_head_outputs: Dict["str", TensorType[..., "num_samples", "features"]],
+        field_output_dict: Dict["str", TensorType[..., "num_samples", "features"]],
         deltas: TensorType[..., "num_samples"],
     ) -> TensorType[..., "out_dim"]:
         """Composite samples along ray and render color image
@@ -62,16 +62,16 @@ class RGB(Renderer):
         Returns:
             TensorType[..., "out_dim"]: Composited RGB ray
         """
-
-        delta_density = deltas * render_head_outputs[self.density_name]
+        delta_density = deltas * field_output_dict[self.density_name][..., 0]
         alphas = 1 - torch.exp(-delta_density)
 
-        transmittance = torch.cat(0, torch.cumsum(delta_density[..., :-1], dim=-1), dim=-1)
+        transmittance = torch.cumsum(delta_density[..., :-1], dim=-1)
+        transmittance = torch.cat([torch.zeros((*transmittance.shape[:1], 1)), transmittance], axis=-1)
         transmittance = torch.exp(-transmittance)  # [..., "num_samples"]
 
         weights = alphas * transmittance  # [..., "num_samples"]
 
-        rgb = torch.sum(weights[..., None] * render_head_outputs[self.rgb_name], dim=-2)
+        rgb = torch.sum(weights[..., None] * field_output_dict[self.rgb_name], dim=-2)
 
         return rgb
 
@@ -83,12 +83,12 @@ class Depth(Renderer):
         Renderer (_type_): _description_
     """
 
-    def required_render_heads(self) -> Set[field_outputs.FieldOutput]:
+    def required_field_outputs(self) -> Set[field_outputs.FieldOutput]:
         return set(field_outputs.DensityFieldOutput)
 
     def forward(
         self,
-        render_head_outputs: Dict["str", TensorType[..., "num_samples", "features"]],
+        field_output_dict: Dict["str", TensorType[..., "num_samples", "features"]],
         deltas: TensorType[..., "num_samples"],
     ) -> TensorType[..., "out_dim"]:
         raise NotImplementedError

tests/nerf/test_renderers.py

@@ -0,0 +1,40 @@
+"""
+Test renderers
+"""
+import pytest
+import torch
+
+from mattport.nerf import renderers
+
+
+def test_rgb_renderer():
+    """Test RGB volumetric rendering"""
+    rgb_name = "rgb"
+    density_name = "density"
+    num_samples = 10
+
+    rgb = torch.ones((3, num_samples, 3))
+    densities = torch.ones((3, num_samples, 1))
+    deltas = torch.ones((3, num_samples))
+
+    rgb_renderer = renderers.RGB(rgb_name, density_name)
+
+    assert rgb_renderer.required_field_outputs is not None
+
+    field_output_dict = {rgb_name: rgb, density_name: densities}
+    out = rgb_renderer(field_output_dict, deltas)
+    assert torch.max(out) > 0.9
+
+    field_output_dict = {rgb_name: rgb * 0, density_name: densities}
+    out = rgb_renderer(field_output_dict, deltas)
+    assert torch.max(out) == pytest.approx(0)
+
+    field_output_dict = {rgb_name: rgb, density_name: densities * 0}
+    out = rgb_renderer(field_output_dict, deltas)
+    assert torch.max(out) == pytest.approx(0)
+
+    # TODO Tancik: Add more precise tests
+
+
+if __name__ == "__main__":
+    test_rgb_renderer()