Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add background_color config to keep consistent with other models #2455

Merged
merged 2 commits into from
Sep 22, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 4 additions & 3 deletions nerfstudio/models/generfacto.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,7 @@
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
Expand All @@ -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)
Expand Down Expand Up @@ -251,7 +252,7 @@ def update_schedule(step):
)

# 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()
Expand Down
4 changes: 2 additions & 2 deletions nerfstudio/models/mipnerf.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@
)
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):
Expand Down Expand Up @@ -81,7 +81,7 @@ def populate_modules(self):
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()

Expand Down
4 changes: 3 additions & 1 deletion nerfstudio/models/tensorf.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -234,7 +236,7 @@ def populate_modules(self):
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()

Expand Down
8 changes: 5 additions & 3 deletions nerfstudio/models/vanilla_nerf.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
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
Expand All @@ -41,7 +41,7 @@
RGBRenderer,
)
from nerfstudio.models.base_model import Model, ModelConfig
from nerfstudio.utils import colormaps, colors, misc
from nerfstudio.utils import colormaps, misc


@dataclass
Expand All @@ -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):
Expand Down Expand Up @@ -110,7 +112,7 @@ def populate_modules(self):
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()

Expand Down
Loading