From cf1794acb16b01796f3d1fcb7501df47a6219df0 Mon Sep 17 00:00:00 2001 From: Terrance Wang Date: Thu, 2 Feb 2023 13:07:14 -0800 Subject: [PATCH] Add lambertian renderer (#1336) * add lambertian shading renderer * adding more functionality * adding to comments * optionally detach normals, fixing naming and typing * Update description --------- Co-authored-by: Matt Tancik Co-authored-by: Matthew Tancik --- nerfstudio/model_components/renderers.py | 34 ++++++++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/nerfstudio/model_components/renderers.py b/nerfstudio/model_components/renderers.py index 24ce46085..1f95de95f 100644 --- a/nerfstudio/model_components/renderers.py +++ b/nerfstudio/model_components/renderers.py @@ -332,3 +332,37 @@ class NormalsRenderer(nn.Module): """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. + + 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