From 4c51bd475534bc9eee3a28a06466c5f5e38dff44 Mon Sep 17 00:00:00 2001 From: Nicholas Pfaff Date: Wed, 5 Jul 2023 18:35:48 +0100 Subject: [PATCH] Clean up typing --- nerfstudio/model_components/renderers.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/nerfstudio/model_components/renderers.py b/nerfstudio/model_components/renderers.py index ca93f3de3..f1e4931e5 100644 --- a/nerfstudio/model_components/renderers.py +++ b/nerfstudio/model_components/renderers.py @@ -39,7 +39,7 @@ from nerfstudio.cameras.rays import RaySamples from nerfstudio.utils import colors from nerfstudio.utils.math import components_from_spherical_harmonics, safe_normalize -BackgroundColor = Union[Literal["random", "last_sample", "black", "white"], Float[Tensor, "3"]] +BackgroundColor = Union[Literal["random", "last_sample", "black", "white"], Float[Tensor, "3"], Float[Tensor, "*bs 3"]] BACKGROUND_COLOR_OVERRIDE: Optional[Float[Tensor, "3"]] = None @@ -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: Union[BackgroundColor, Tensor] = "random", + background_color: BackgroundColor = "random", ray_indices: Optional[Int[Tensor, "num_samples"]] = None, num_rays: Optional[int] = None, ) -> Float[Tensor, "*bs 3"]: @@ -108,8 +108,8 @@ class RGBRenderer(nn.Module): @classmethod def get_background_color( - cls, rgb: Float[Tensor, "*bs num_samples 3"], background_color: Union[BackgroundColor, Tensor] - ) -> Tensor: + cls, rgb: Float[Tensor, "*bs num_samples 3"], background_color: BackgroundColor + ) -> Union[Float[Tensor, "3"], Float[Tensor, "*bs 3"]]: """Returns the RGB background color for a specified background color. Args: @@ -154,11 +154,11 @@ class RGBRenderer(nn.Module): def blend_background_for_metric_computation( self, - gt_image: torch.Tensor, - predicted_rgb: torch.Tensor, - predicted_opacity: torch.Tensor, - background_color: torch.Tensor = colors.BLACK, - ) -> Tuple[torch.Tensor, torch.Tensor]: + gt_image: Tensor, + predicted_rgb: Tensor, + predicted_opacity: Tensor, + background_color: Float[Tensor, "3"] = colors.BLACK, + ) -> Tuple[Tensor, Tensor]: """Blends a background color into the ground truth and predicted image for metric computation if the current background color is "random". @@ -190,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: Optional[Union[BackgroundColor, Tensor]] = None, + background_color: Optional[BackgroundColor] = None, ) -> Float[Tensor, "*bs 3"]: """Composite samples along ray and render color image