From 2e79c93cd2562e8bb8c97abf5104cd13141ee804 Mon Sep 17 00:00:00 2001 From: Matthew Tancik Date: Thu, 2 Feb 2023 18:06:11 -0800 Subject: [PATCH 1/4] Fix normals visualization (#1341) --- nerfstudio/model_components/renderers.py | 13 +++++++++++-- nerfstudio/models/nerfacto.py | 6 ++++-- 2 files changed, 15 insertions(+), 4 deletions(-) diff --git a/nerfstudio/model_components/renderers.py b/nerfstudio/model_components/renderers.py index 1f95de95f..6a8e2a5f9 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,9 +328,18 @@ 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.""" + """Calculate normals along the ray. + + Args: + normals: Normals for each sample. + weights: Weights of each sample. + normalize: Normalize normals. + """ n = torch.sum(weights * normals, dim=-2) + if normalize: + n = safe_normalize(n) return n diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index c996e20b3..91b044884 100644 --- a/nerfstudio/models/nerfacto.py +++ b/nerfstudio/models/nerfacto.py @@ -250,8 +250,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"] = (normals + 1) / 2 + outputs["pred_normals"] = (pred_normals + 1) / 2 # These use a lot of GPU memory, so we avoid storing them for eval. if self.training: From a0976d1ff82a4cd3f34860071fefaa541b589357 Mon Sep 17 00:00:00 2001 From: Jeffrey Hu <102542936+jeffreyhuparallel@users.noreply.github.com> Date: Thu, 2 Feb 2023 18:22:55 -0800 Subject: [PATCH 2/4] Fix depth-supervised nerf when using z-buffer depth. (#1340) * fix z buffer depth * remove script --- nerfstudio/model_components/losses.py | 2 +- nerfstudio/models/depth_nerfacto.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) 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/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( From 59d7e99c804de7588b9ffd5d993f1c667a2a5d23 Mon Sep 17 00:00:00 2001 From: Matthew Tancik Date: Fri, 3 Feb 2023 11:21:54 -0800 Subject: [PATCH 3/4] Create shaders (#1344) --- nerfstudio/model_components/renderers.py | 34 ----------- nerfstudio/model_components/shaders.py | 77 ++++++++++++++++++++++++ nerfstudio/models/nerfacto.py | 15 ++--- 3 files changed, 83 insertions(+), 43 deletions(-) create mode 100644 nerfstudio/model_components/shaders.py diff --git a/nerfstudio/model_components/renderers.py b/nerfstudio/model_components/renderers.py index 6a8e2a5f9..d13c29233 100644 --- a/nerfstudio/model_components/renderers.py +++ b/nerfstudio/model_components/renderers.py @@ -341,37 +341,3 @@ class NormalsRenderer(nn.Module): if normalize: n = safe_normalize(n) 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. - - 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 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/nerfacto.py b/nerfstudio/models/nerfacto.py index 91b044884..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() @@ -252,9 +256,8 @@ class NerfactoModel(Model): if self.config.predict_normals: 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"] = (normals + 1) / 2 - outputs["pred_normals"] = (pred_normals + 1) / 2 - + 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 @@ -336,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( From a1ea432bf71e6cb69a5553d6b64be562d888644a Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Fri, 3 Feb 2023 11:32:12 -0800 Subject: [PATCH 4/4] Use shaders --- nerfstudio/generative/stable_diffusion.py | 2 +- nerfstudio/models/dreamfusion.py | 33 +++++++---------------- 2 files changed, 11 insertions(+), 24 deletions(-) diff --git a/nerfstudio/generative/stable_diffusion.py b/nerfstudio/generative/stable_diffusion.py index eb16710f0..1816da32d 100644 --- a/nerfstudio/generative/stable_diffusion.py +++ b/nerfstudio/generative/stable_diffusion.py @@ -92,7 +92,7 @@ class StableDiffusion(nn.Module): unet_traced_filename = Path(appdirs.user_data_dir("nerfstudio")) / "unet_traced.pt" if unet_traced_filename.exists(): CONSOLE.print("Loading traced UNet.") - unet_traced = torch.jit.load("unet_traced.pt") + unet_traced = torch.jit.load(unet_traced_filename) class TracedUNet(torch.nn.Module): """Jitted version of UNet""" 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