From 2fd8b59d20e73f55106012623d188b227a976df5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E3=83=95=E3=83=A9=E3=83=8F=E3=83=86=E3=82=A3=E3=80=80?= =?UTF-8?q?=E4=BB=81?= Date: Fri, 22 Sep 2023 23:40:20 +0900 Subject: [PATCH] Add `background_color` config to keep consistent with other models (#2455) * add backgorund_color config * removed unused color import --- nerfstudio/models/generfacto.py | 7 ++++--- nerfstudio/models/mipnerf.py | 4 ++-- nerfstudio/models/tensorf.py | 4 +++- nerfstudio/models/vanilla_nerf.py | 8 +++++--- 4 files changed, 14 insertions(+), 9 deletions(-) diff --git a/nerfstudio/models/generfacto.py b/nerfstudio/models/generfacto.py index bf34eb335..c07f904c7 100644 --- a/nerfstudio/models/generfacto.py +++ b/nerfstudio/models/generfacto.py @@ -57,7 +57,7 @@ from nerfstudio.model_components.renderers import ( 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 +from nerfstudio.utils import colormaps, math, misc @dataclass @@ -68,7 +68,8 @@ class GenerfactoModelConfig(ModelConfig): """target class to instantiate""" prompt: str = "a high quality photo of a ripe pineapple" """prompt for stable dreamfusion""" - + background_color: Literal["random", "last_sample", "black", "white"] = "white" + """Whether to randomize the background color.""" orientation_loss_mult: Tuple[float, float] = (0.001, 10.0) """Orientation loss multipier on computed normals.""" orientation_loss_mult_range: Tuple[int, int] = (0, 15000) @@ -251,7 +252,7 @@ class GenerfactoModel(Model): ) # renderers - self.renderer_rgb = RGBRenderer(background_color=colors.WHITE) + self.renderer_rgb = RGBRenderer(background_color=self.config.background_color) self.renderer_accumulation = AccumulationRenderer() self.renderer_depth = DepthRenderer() self.renderer_normals = NormalsRenderer() diff --git a/nerfstudio/models/mipnerf.py b/nerfstudio/models/mipnerf.py index a3a783461..67ca63251 100644 --- a/nerfstudio/models/mipnerf.py +++ b/nerfstudio/models/mipnerf.py @@ -38,7 +38,7 @@ from nerfstudio.model_components.renderers import ( ) from nerfstudio.models.base_model import Model from nerfstudio.models.vanilla_nerf import VanillaModelConfig -from nerfstudio.utils import colormaps, colors, misc +from nerfstudio.utils import colormaps, misc class MipNerfModel(Model): @@ -81,7 +81,7 @@ class MipNerfModel(Model): self.sampler_pdf = PDFSampler(num_samples=self.config.num_importance_samples, include_original=False) # renderers - self.renderer_rgb = RGBRenderer(background_color=colors.WHITE) + self.renderer_rgb = RGBRenderer(background_color=self.config.background_color) self.renderer_accumulation = AccumulationRenderer() self.renderer_depth = DepthRenderer() diff --git a/nerfstudio/models/tensorf.py b/nerfstudio/models/tensorf.py index a2cd9901d..a0380d42e 100644 --- a/nerfstudio/models/tensorf.py +++ b/nerfstudio/models/tensorf.py @@ -89,6 +89,8 @@ class TensoRFModelConfig(ModelConfig): tensorf_encoding: Literal["triplane", "vm", "cp"] = "vm" regularization: Literal["none", "l1", "tv"] = "l1" """Regularization method used in tensorf paper""" + background_color: Literal["random", "last_sample", "black", "white"] = "white" + """Whether to randomize the background color.""" class TensoRFModel(Model): @@ -234,7 +236,7 @@ class TensoRFModel(Model): self.sampler_pdf = PDFSampler(num_samples=self.config.num_samples, single_jitter=True, include_original=False) # renderers - self.renderer_rgb = RGBRenderer(background_color=colors.WHITE) + self.renderer_rgb = RGBRenderer(background_color=self.config.background_color) self.renderer_accumulation = AccumulationRenderer() self.renderer_depth = DepthRenderer() diff --git a/nerfstudio/models/vanilla_nerf.py b/nerfstudio/models/vanilla_nerf.py index ea5344fb8..e95c8dfa4 100644 --- a/nerfstudio/models/vanilla_nerf.py +++ b/nerfstudio/models/vanilla_nerf.py @@ -19,7 +19,7 @@ Implementation of vanilla nerf. from __future__ import annotations from dataclasses import dataclass, field -from typing import Any, Dict, List, Tuple, Type +from typing import Any, Dict, List, Tuple, Type, Literal import torch from torch.nn import Parameter @@ -41,7 +41,7 @@ from nerfstudio.model_components.renderers import ( RGBRenderer, ) from nerfstudio.models.base_model import Model, ModelConfig -from nerfstudio.utils import colormaps, colors, misc +from nerfstudio.utils import colormaps, misc @dataclass @@ -58,6 +58,8 @@ class VanillaModelConfig(ModelConfig): """Specifies whether or not to include ray warping based on time.""" temporal_distortion_params: Dict[str, Any] = to_immutable_dict({"kind": TemporalDistortionKind.DNERF}) """Parameters to instantiate temporal distortion with""" + background_color: Literal["random", "last_sample", "black", "white"] = "white" + """Whether to randomize the background color.""" class NeRFModel(Model): @@ -110,7 +112,7 @@ class NeRFModel(Model): self.sampler_pdf = PDFSampler(num_samples=self.config.num_importance_samples) # renderers - self.renderer_rgb = RGBRenderer(background_color=colors.WHITE) + self.renderer_rgb = RGBRenderer(background_color=self.config.background_color) self.renderer_accumulation = AccumulationRenderer() self.renderer_depth = DepthRenderer()