From a96088054a4c04d5ae266e6eab66fa8d1a9c0014 Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Fri, 15 Apr 2022 16:58:28 -0700 Subject: [PATCH] Fix renderer, add test --- .github/workflows/code_checks.yml | 2 +- .vscode/settings.json | 3 ++- mattport/nerf/renderers.py | 20 ++++++++-------- tests/nerf/test_renderers.py | 40 +++++++++++++++++++++++++++++++ 4 files changed, 53 insertions(+), 12 deletions(-) create mode 100644 tests/nerf/test_renderers.py diff --git a/.github/workflows/code_checks.yml b/.github/workflows/code_checks.yml index 9dd89fe5a..b073d6757 100644 --- a/.github/workflows/code_checks.yml +++ b/.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 diff --git a/.vscode/settings.json b/.vscode/settings.json index 3ce315795..577383dfe 100644 --- a/.vscode/settings.json +++ b/.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" ], } \ No newline at end of file diff --git a/mattport/nerf/renderers.py b/mattport/nerf/renderers.py index d319ea74e..805920db2 100644 --- a/mattport/nerf/renderers.py +++ b/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 diff --git a/tests/nerf/test_renderers.py b/tests/nerf/test_renderers.py new file mode 100644 index 000000000..e6897dad9 --- /dev/null +++ b/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()