diff --git a/nerfstudio/model_components/losses.py b/nerfstudio/model_components/losses.py index 7fcfe0419..6a981d664 100644 --- a/nerfstudio/model_components/losses.py +++ b/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: diff --git a/nerfstudio/model_components/renderers.py b/nerfstudio/model_components/renderers.py index 1f95de95f..d13c29233 100644 --- a/nerfstudio/model_components/renderers.py +++ b/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 diff --git a/nerfstudio/model_components/shaders.py b/nerfstudio/model_components/shaders.py new file mode 100644 index 000000000..e71388e27 --- /dev/null +++ b/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 diff --git a/nerfstudio/models/depth_nerfacto.py b/nerfstudio/models/depth_nerfacto.py index da4c59002..ac0768fdf 100644 --- a/nerfstudio/models/depth_nerfacto.py +++ b/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( diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py index e78600a03..9588aa7b1 100644 --- a/nerfstudio/models/dreamfusion.py +++ b/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 diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index c996e20b3..7442bfe8b 100644 --- a/nerfstudio/models/nerfacto.py +++ b/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(