From 32900a6e5c53dd6824f4c8009afb9670d71fb227 Mon Sep 17 00:00:00 2001 From: Nicholas Pfaff Date: Wed, 5 Jul 2023 18:30:24 +0100 Subject: [PATCH] Fix typing error --- nerfstudio/model_components/renderers.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/nerfstudio/model_components/renderers.py b/nerfstudio/model_components/renderers.py index ac7a19e5f..ca93f3de3 100644 --- a/nerfstudio/model_components/renderers.py +++ b/nerfstudio/model_components/renderers.py @@ -71,7 +71,7 @@ class RGBRenderer(nn.Module): cls, rgb: Float[Tensor, "*bs num_samples 3"], weights: Float[Tensor, "*bs num_samples 1"], - background_color: BackgroundColor = "random", + background_color: Union[BackgroundColor, Tensor] = "random", ray_indices: Optional[Int[Tensor, "num_samples"]] = None, num_rays: Optional[int] = None, ) -> Float[Tensor, "*bs 3"]: @@ -107,7 +107,9 @@ class RGBRenderer(nn.Module): return comp_rgb @classmethod - def get_background_color(cls, rgb: Float[Tensor, "*bs num_samples 3"], background_color: BackgroundColor) -> Tensor: + def get_background_color( + cls, rgb: Float[Tensor, "*bs num_samples 3"], background_color: Union[BackgroundColor, Tensor] + ) -> Tensor: """Returns the RGB background color for a specified background color. Args: @@ -188,7 +190,7 @@ class RGBRenderer(nn.Module): weights: Float[Tensor, "*bs num_samples 1"], ray_indices: Optional[Int[Tensor, "num_samples"]] = None, num_rays: Optional[int] = None, - background_color: Union[Float[Tensor, "3"], BackgroundColor, None] = None, + background_color: Optional[Union[BackgroundColor, Tensor]] = None, ) -> Float[Tensor, "*bs 3"]: """Composite samples along ray and render color image