Tune proposal network (#638)

nerfstudio/configs/model_configs.py

@@ -62,7 +62,7 @@ model_configs["nerfacto"] = Config(
         datamanager=VanillaDataManagerConfig(
             dataparser=NerfstudioDataParserConfig(), train_num_rays_per_batch=4096, eval_num_rays_per_batch=8192
         ),
-        model=NerfactoModelConfig(eval_num_rays_per_chunk=1 << 16),
+        model=NerfactoModelConfig(eval_num_rays_per_chunk=1 << 14),
     ),
     optimizers={
         "proposal_networks": {
@@ -74,7 +74,7 @@ model_configs["nerfacto"] = Config(
             "scheduler": None,
         },
     },
-    viewer=ViewerConfig(num_rays_per_chunk=1 << 16),
+    viewer=ViewerConfig(num_rays_per_chunk=1 << 14),
     vis="viewer",
 )
 

nerfstudio/model_components/ray_samplers.py

@@ -59,6 +59,7 @@ class SpacedSampler(Sampler):
         spacing_fn: Function that dictates sample spacing (ie `lambda x : x` is uniform).
         spacing_fn_inv: The inverse of spacing_fn.
         train_stratified: Use stratified sampling during training. Defults to True
+        single_jitter: Use a same random jitter for all samples along a ray. Defaults to False
     """
 
     def __init__(
@@ -131,6 +132,7 @@ class UniformSampler(SpacedSampler):
     Args:
         num_samples: Number of samples per ray
         train_stratified: Use stratified sampling during training. Defults to True
+        single_jitter: Use a same random jitter for all samples along a ray. Defaults to False
     """
 
     def __init__(
@@ -154,6 +156,7 @@ class LinearDisparitySampler(SpacedSampler):
     Args:
         num_samples: Number of samples per ray
         train_stratified: Use stratified sampling during training. Defults to True
+        single_jitter: Use a same random jitter for all samples along a ray. Defaults to False
     """
 
     def __init__(
@@ -225,6 +228,7 @@ class UniformLinDispPiecewiseSampler(SpacedSampler):
     Args:
         num_samples: Number of samples per ray
         train_stratified: Use stratified sampling during training. Defults to True
+        single_jitter: Use a same random jitter for all samples along a ray. Defaults to False
     """
 
     def __init__(
@@ -248,6 +252,7 @@ class PDFSampler(Sampler):
     Args:
         num_samples: Number of samples per ray
         train_stratified: Randomize location within each bin during training.
+        single_jitter: Use a same random jitter for all samples along a ray. Defaults to False
         include_original: Add original samples to ray.
         histogram_padding: Amount to weights prior to computing PDF.
     """
@@ -256,6 +261,7 @@ class PDFSampler(Sampler):
         self,
         num_samples: Optional[int] = None,
         train_stratified: bool = True,
+        single_jitter: bool = False,
         include_original: bool = True,
         histogram_padding: float = 0.01,
     ) -> None:
@@ -263,6 +269,7 @@ class PDFSampler(Sampler):
         self.train_stratified = train_stratified
         self.include_original = include_original
         self.histogram_padding = histogram_padding
+        self.single_jitter = single_jitter
 
     def generate_ray_samples(
         self,
@@ -308,7 +315,12 @@ class PDFSampler(Sampler):
             # Stratified samples between 0 and 1
             u = torch.linspace(0.0, 1.0 - (1.0 / num_bins), steps=num_bins, device=cdf.device)
             u = u.expand(size=(*cdf.shape[:-1], num_bins))
-            u = u + torch.rand(size=(*cdf.shape[:-1], num_bins), device=cdf.device) / num_bins
+            # u = u + torch.rand(size=(*cdf.shape[:-1], num_bins), device=cdf.device) / num_bins
+            if self.single_jitter:
+                rand = torch.rand((*cdf.shape[:-1], 1), device=cdf.device) / num_bins
+            else:
+                rand = torch.rand((*cdf.shape[:-1], num_samples + 1), device=cdf.device) / num_bins
+            u = u + rand
         else:
             # Uniform samples between 0 and 1
             u = torch.linspace(0.0, 1.0 - (1.0 / num_bins), steps=num_bins, device=cdf.device)
@@ -505,7 +517,7 @@ class ProposalNetworkSampler(Sampler):
 
         # samplers
         self.initial_sampler = UniformLinDispPiecewiseSampler(single_jitter=single_jitter)
-        self.pdf_sampler = PDFSampler(include_original=False)
+        self.pdf_sampler = PDFSampler(include_original=False, single_jitter=single_jitter)
 
         self._anneal = 1.0
 

nerfstudio/model_components/renderers.py

@@ -27,7 +27,7 @@ Example:
 
 """
 import math
-from typing import Literal, Optional
+from typing import Literal, Optional, Union
 
 import nerfacc
 import torch
@@ -45,7 +45,7 @@ class RGBRenderer(nn.Module):
         background_color: Background color as RGB. Uses random colors if None.
     """
 
-    def __init__(self, background_color: Optional[TensorType[3]] = None) -> None:
+    def __init__(self, background_color: Union[Literal["random", "last_sample"], TensorType[3]] = "random") -> None:
         super().__init__()
         self.background_color = background_color
 
@@ -54,7 +54,7 @@ class RGBRenderer(nn.Module):
         cls,
         rgb: TensorType["bs":..., "num_samples", 3],
         weights: TensorType["bs":..., "num_samples", 1],
-        background_color: Optional[TensorType[3]] = None,
+        background_color: Union[Literal["random", "last_sample"], TensorType[3]] = "random",
         ray_indices: Optional[TensorType["num_samples"]] = None,
         num_rays: Optional[int] = None,
     ) -> TensorType["bs":..., 3]:
@@ -63,7 +63,7 @@ class RGBRenderer(nn.Module):
         Args:
             rgb: RGB for each sample
             weights: Weights for each sample
-            background_color: Background color as RGB. Uses random colors if None.
+            background_color: Background color as RGB.
             ray_indices: Ray index for each sample, used when samples are packed.
             num_rays: Number of rays, used when samples are packed.
 
@@ -72,18 +72,23 @@ class RGBRenderer(nn.Module):
         """
         if ray_indices is not None and num_rays is not None:
             # Necessary for packed samples from volumetric ray sampler
-            rgb = nerfacc.accumulate_along_rays(weights, ray_indices, rgb, num_rays)
+            if background_color == "last_sample":
+                raise NotImplementedError("Background color 'last_sample' not implemented for packed samples.")
+            comp_rgb = nerfacc.accumulate_along_rays(weights, ray_indices, rgb, num_rays)
             accumulated_weight = nerfacc.accumulate_along_rays(weights, ray_indices, None, num_rays)
         else:
-            rgb = torch.sum(weights * rgb, dim=-2)
+            comp_rgb = torch.sum(weights * rgb, dim=-2)
             accumulated_weight = torch.sum(weights, dim=-2)
 
-        if background_color is None:
-            background_color = torch.rand_like(rgb).to(rgb.device)
+        if background_color == "last_sample":
+            background_color = rgb[..., -1, :]
+        if background_color == "random":
+            background_color = torch.rand_like(comp_rgb).to(rgb.device)
 
-        rgb = rgb + background_color.to(weights.device) * (1.0 - accumulated_weight)
+        assert isinstance(background_color, torch.Tensor)
+        comp_rgb = comp_rgb + background_color.to(weights.device) * (1.0 - accumulated_weight)
 
-        return rgb
+        return comp_rgb
 
     def forward(
         self,
@@ -121,7 +126,9 @@ class SHRenderer(nn.Module):
     """
 
     def __init__(
-        self, background_color: Optional[TensorType[3]] = None, activation: Optional[nn.Module] = nn.Sigmoid()
+        self,
+        background_color: Union[Literal["random", "last_sample"], TensorType[3]] = "random",
+        activation: Optional[nn.Module] = nn.Sigmoid(),
     ) -> None:
         super().__init__()
         self.background_color = background_color

nerfstudio/models/instant_ngp.py

@@ -122,7 +122,7 @@ class NGPModel(Model):
         )
 
         # renderers
-        background_color = None if self.config.randomize_background else colors.WHITE
+        background_color = "random" if self.config.randomize_background else colors.WHITE
         self.renderer_rgb = RGBRenderer(background_color=background_color)
         self.renderer_accumulation = AccumulationRenderer()
         self.renderer_depth = DepthRenderer()

nerfstudio/models/nerfacto.py

@@ -19,7 +19,7 @@ NeRF implementation that combines many recent advancements.
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from typing import Dict, List, Tuple, Type
+from typing import Dict, List, Literal, Tuple, Type
 
 import numpy as np
 import torch
@@ -47,7 +47,7 @@ from nerfstudio.model_components.renderers import (
 )
 from nerfstudio.model_components.scene_colliders import NearFarCollider
 from nerfstudio.models.base_model import Model, ModelConfig
-from nerfstudio.utils import colormaps, colors
+from nerfstudio.utils import colormaps
 
 
 @dataclass
@@ -59,7 +59,7 @@ class NerfactoModelConfig(ModelConfig):
     """How far along the ray to start sampling."""
     far_plane: float = 1000.0
     """How far along the ray to stop sampling."""
-    randomize_background: bool = True
+    background_color: Literal["background", "last_sample"] = "last_sample"
     """Whether to randomize the background color."""
     num_proposal_samples_per_ray: Tuple[int] = (64,)
     """Number of samples per ray for the proposal network."""
@@ -71,10 +71,8 @@ class NerfactoModelConfig(ModelConfig):
     """Use the same proposal network. Otherwise use different ones."""
     interlevel_loss_mult: float = 1.0
     """Proposal loss multiplier."""
-    distortion_loss_mult: float = 0.01
+    distortion_loss_mult: float = 0.002
     """Distortion loss multiplier."""
-    use_appearance_conditioning: bool = True
-    """Whether to use appearance conditioning."""
     use_proposal_weight_anneal: bool = True
     """Whether to use proposal weight annealing."""
     use_average_appearance_embedding: bool = True
@@ -83,8 +81,8 @@ class NerfactoModelConfig(ModelConfig):
     """Slope of the annealing function for the proposal weights."""
     proposal_weights_anneal_max_num_iters: int = 1000
     """Max num iterations for the annealing function."""
-    use_single_jitter: bool = False
-    """Whether use single jitter or not for first proposal network."""
+    use_single_jitter: bool = True
+    """Whether use single jitter or not for the proposal networks."""
 
 
 class NerfactoModel(Model):
@@ -103,15 +101,12 @@ class NerfactoModel(Model):
         scene_contraction = SceneContraction(order=float("inf"))
 
         # Fields
-        if self.config.use_appearance_conditioning:
-            self.field = TCNNNerfactoField(
-                self.scene_box.aabb,
-                spatial_distortion=scene_contraction,
-                num_images=self.num_train_data,
-                use_average_appearance_embedding=self.config.use_average_appearance_embedding,
-            )
-        else:
-            raise NotImplementedError("Only appearance conditioning is supported.")
+        self.field = TCNNNerfactoField(
+            self.scene_box.aabb,
+            spatial_distortion=scene_contraction,
+            num_images=self.num_train_data,
+            use_average_appearance_embedding=self.config.use_average_appearance_embedding,
+        )
 
         # Build the proposal network(s)
         self.proposal_networks = torch.nn.ModuleList()
@@ -137,8 +132,7 @@ class NerfactoModel(Model):
         )
 
         # renderers
-        background_color = None if self.config.randomize_background else colors.WHITE
-        self.renderer_rgb = RGBRenderer(background_color=background_color)
+        self.renderer_rgb = RGBRenderer(background_color=self.config.background_color)
         self.renderer_accumulation = AccumulationRenderer()
         self.renderer_depth = DepthRenderer()
 

nerfstudio/viewer/server/viewer_utils.py

@@ -537,8 +537,12 @@ class ViewerState:
         image_height = (num_vis_rays / aspect_ratio) ** 0.5
         image_height = int(round(image_height, -1))
         image_height = min(self.max_resolution, image_height)
+        image_width = int(image_height * aspect_ratio)
+        if image_width > self.max_resolution:
+            image_width = self.max_resolution
+            image_height = int(image_width / aspect_ratio)
 
-        return image_height, int(image_height * aspect_ratio)
+        return image_height, image_width
 
     def _process_invalid_output(self, output_type: str) -> str:
         """Check to see whether we are in the corner case of RGB; if still invalid, throw error