Merge branch 'generative' of https://github.com/plenoptix/nerfactory into generative

nerfstudio/model_components/losses.py

@@ -291,7 +291,7 @@ def depth_loss(
         Depth loss scalar.
     """
     if not is_euclidean:
-        termination_depth = termination_depth / directions_norm
+        termination_depth = termination_depth * directions_norm
     steps = (ray_samples.frustums.starts + ray_samples.frustums.ends) / 2
 
     if depth_loss_type == DephtLossType.DS_NERF:

nerfstudio/model_components/renderers.py

@@ -37,7 +37,7 @@ from torchtyping import TensorType
 from typing_extensions import Literal
 
 from nerfstudio.cameras.rays import RaySamples
-from nerfstudio.utils.math import components_from_spherical_harmonics
+from nerfstudio.utils.math import components_from_spherical_harmonics, safe_normalize
 
 BACKGROUND_COLOR_OVERRIDE: Optional[TensorType[3]] = None
 
@@ -328,41 +328,16 @@ class NormalsRenderer(nn.Module):
         cls,
         normals: TensorType["bs":..., "num_samples", 3],
         weights: TensorType["bs":..., "num_samples", 1],
+        normalize: bool = True,
     ) -> TensorType["bs":..., 3]:
-        """Calculate normals along the ray."""
-        n = torch.sum(weights * normals, dim=-2)
-        return n
-
-
-class LambertianShadingRenderer(nn.Module):
-    """Calculate Lambertian shading."""
-
-    @classmethod
-    def forward(
-        cls,
-        rgb: TensorType["bs":..., 3],
-        normals: TensorType["bs":..., 3],
-        light_direction: TensorType["bs":..., 3],
-        shading_weight: float = 1.0,
-        detach_normals=True,
-    ):
-        """Calculate Lambertian shading.
+        """Calculate normals along the ray.
 
         Args:
-            rgb: Accumulated rgb along a ray.
-            normals: Accumulated normals along a ray.
-            light_direction: Direction of light source.
-            shading_weight: Lambertian shading (1.0) vs. ambient lighting (0.0) ratio
-            detach_normals: Detach normals from the computation graph when computing shading.
-
-        Returns:
-            Textureless Lambertian shading, Lambertian shading
+            normals: Normals for each sample.
+            weights: Weights of each sample.
+            normalize: Normalize normals.
         """
-        if detach_normals:
-            normals = normals.detach()
-
-        lambertian = (1 - shading_weight) + shading_weight * (normals @ light_direction).clamp(min=0)
-        shaded = lambertian.unsqueeze(-1).repeat(1, 3)
-        shaded_albedo = rgb * lambertian.unsqueeze(-1)
-
-        return shaded, shaded_albedo
+        n = torch.sum(weights * normals, dim=-2)
+        if normalize:
+            n = safe_normalize(n)
+        return n

nerfstudio/model_components/shaders.py

@@ -0,0 +1,77 @@
+# Copyright 2022 The Nerfstudio Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Shaders for rendering."""
+from typing import Optional
+
+from torch import nn
+from torchtyping import TensorType
+
+
+class LambertianShader(nn.Module):
+    """Calculate Lambertian shading."""
+
+    @classmethod
+    def forward(
+        cls,
+        rgb: TensorType["bs":..., 3],
+        normals: TensorType["bs":..., 3],
+        light_direction: TensorType["bs":..., 3],
+        shading_weight: float = 1.0,
+        detach_normals=True,
+    ):
+        """Calculate Lambertian shading.
+
+        Args:
+            rgb: Accumulated rgb along a ray.
+            normals: Accumulated normals along a ray.
+            light_direction: Direction of light source.
+            shading_weight: Lambertian shading (1.0) vs. ambient lighting (0.0) ratio
+            detach_normals: Detach normals from the computation graph when computing shading.
+
+        Returns:
+            Textureless Lambertian shading, Lambertian shading
+        """
+        if detach_normals:
+            normals = normals.detach()
+
+        lambertian = (1 - shading_weight) + shading_weight * (normals @ light_direction).clamp(min=0)
+        shaded = lambertian.unsqueeze(-1).repeat(1, 3)
+        shaded_albedo = rgb * lambertian.unsqueeze(-1)
+
+        return shaded, shaded_albedo
+
+
+class NormalsShader(nn.Module):
+    """Calculate shading for normals."""
+
+    @classmethod
+    def forward(
+        cls,
+        normals: TensorType["bs":..., 3],
+        weights: Optional[TensorType["bs":..., 1]] = None,
+    ):
+        """Applies a rainbow colormap to the normals.
+
+        Args:
+            normals: Normalized 3D vectors.
+            weights: Optional weights to scale to the normal colors. (Can be used for masking)
+
+        Returns:
+            Colored normals
+        """
+        normals = (normals + 1) / 2
+        if weights is not None:
+            normals = normals * weights
+        return normals

nerfstudio/models/depth_nerfacto.py

@@ -108,7 +108,7 @@ class DepthNerfactoModel(NerfactoModel):
         metrics, images = super().get_image_metrics_and_images(outputs, batch)
         ground_truth_depth = batch["depth_image"]
         if not self.config.is_euclidean_depth:
-            ground_truth_depth = ground_truth_depth / outputs["directions_norm"]
+            ground_truth_depth = ground_truth_depth * outputs["directions_norm"]
 
         ground_truth_depth_colormap = colormaps.apply_depth_colormap(ground_truth_depth)
         predicted_depth_colormap = colormaps.apply_depth_colormap(

nerfstudio/models/dreamfusion.py

@@ -52,11 +52,11 @@ from nerfstudio.model_components.ray_samplers import (
 from nerfstudio.model_components.renderers import (
     AccumulationRenderer,
     DepthRenderer,
-    LambertianShadingRenderer,
     NormalsRenderer,
     RGBRenderer,
 )
 from nerfstudio.model_components.scene_colliders import AABBBoxCollider, SphereCollider
+from nerfstudio.model_components.shaders import LambertianShader, NormalsShader
 from nerfstudio.models.base_model import Model, ModelConfig
 from nerfstudio.utils import colormaps, colors, math, misc
 
@@ -183,7 +183,10 @@ class DreamFusionModel(Model):
         self.renderer_accumulation = AccumulationRenderer()
         self.renderer_depth = DepthRenderer()
         self.renderer_normals = NormalsRenderer()
-        self.renderer_lambertian = LambertianShadingRenderer()
+
+        # shaders
+        self.shader_lambertian = LambertianShader()
+        self.shader_normals = NormalsShader()
 
         # losses
         self.rgb_loss = MSELoss()
@@ -340,12 +343,12 @@ class DreamFusionModel(Model):
         else:
             shading_weight = 0.0
 
-        shaded, shaded_albedo = self.renderer_lambertian(
+        shaded, shaded_albedo = self.shader_lambertian(
             rgb=rgb, normals=pred_normals, light_direction=light_d, shading_weight=shading_weight, detach_normals=True
         )
 
-        outputs["normals"] = ((normals + 1.0) / 2.0) * accum_mask
-        outputs["pred_normals"] = ((pred_normals + 1.0) / 2.0) * accum_mask
+        outputs["normals"] = self.shader_normals(normals, weights=accum_mask)
+        outputs["pred_normals"] = self.shader_normals(pred_normals, weights=accum_mask)
         outputs["shaded"] = accum_mask * shaded + torch.ones_like(shaded) * accum_mask_inv
         outputs["shaded_albedo"] = shaded_albedo
         outputs["render"] = accum_mask * rgb + background
@@ -409,28 +412,12 @@ class DreamFusionModel(Model):
     def get_image_metrics_and_images(
         self, outputs: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor]
     ) -> Tuple[Dict[str, float], Dict[str, torch.Tensor]]:
-        image = batch["image"].to(outputs["rgb"].device)
-        rgb = outputs["rgb"]
         acc = colormaps.apply_colormap(outputs["accumulation"])
         depth = colormaps.apply_depth_colormap(
             outputs["depth"],
             accumulation=outputs["accumulation"],
         )
 
-        combined_rgb = torch.cat([image, rgb], dim=1)
-
-        # Switch images from [H, W, C] to [1, C, H, W] for metrics computations
-        image = torch.moveaxis(image, -1, 0)[None, ...]
-        rgb = torch.moveaxis(rgb, -1, 0)[None, ...]
-
-        psnr = self.psnr(image, rgb)
-        ssim = self.ssim(image, rgb)
-        lpips = self.lpips(image, rgb)
-
-        metrics_dict = {
-            "psnr": float(psnr.item()),
-            "ssim": float(ssim.item()),
-            "lpips": float(lpips.item()),
-        }
-        images_dict = {"img": combined_rgb, "accumulation": acc, "depth": depth}
+        metrics_dict = {}
+        images_dict = {"img": outputs["rgb"], "accumulation": acc, "depth": depth}
         return metrics_dict, images_dict

nerfstudio/models/nerfacto.py

@@ -54,6 +54,7 @@ from nerfstudio.model_components.renderers import (
     RGBRenderer,
 )
 from nerfstudio.model_components.scene_colliders import NearFarCollider
+from nerfstudio.model_components.shaders import NormalsShader
 from nerfstudio.models.base_model import Model, ModelConfig
 from nerfstudio.utils import colormaps
 
@@ -187,6 +188,9 @@ class NerfactoModel(Model):
         self.renderer_depth = DepthRenderer()
         self.renderer_normals = NormalsRenderer()
 
+        # shaders
+        self.normals_shader = NormalsShader()
+
         # losses
         self.rgb_loss = MSELoss()
 
@@ -250,9 +254,10 @@ class NerfactoModel(Model):
         }
 
         if self.config.predict_normals:
-            outputs["normals"] = self.renderer_normals(normals=field_outputs[FieldHeadNames.NORMALS], weights=weights)
-            outputs["pred_normals"] = self.renderer_normals(field_outputs[FieldHeadNames.PRED_NORMALS], weights=weights)
-
+            normals = self.renderer_normals(normals=field_outputs[FieldHeadNames.NORMALS], weights=weights)
+            pred_normals = self.renderer_normals(field_outputs[FieldHeadNames.PRED_NORMALS], weights=weights)
+            outputs["normals"] = self.normals_shader(normals)
+            outputs["pred_normals"] = self.normals_shader(pred_normals)
         # These use a lot of GPU memory, so we avoid storing them for eval.
         if self.training:
             outputs["weights_list"] = weights_list
@@ -334,12 +339,6 @@ class NerfactoModel(Model):
 
         images_dict = {"img": combined_rgb, "accumulation": combined_acc, "depth": combined_depth}
 
-        # normals to RGB for visualization. TODO: use a colormap
-        if "normals" in outputs:
-            images_dict["normals"] = (outputs["normals"] + 1.0) / 2.0
-        if "pred_normals" in outputs:
-            images_dict["pred_normals"] = (outputs["pred_normals"] + 1.0) / 2.0
-
         for i in range(self.config.num_proposal_iterations):
             key = f"prop_depth_{i}"
             prop_depth_i = colormaps.apply_depth_colormap(