Merge branch 'generative' of https://github.com/plenoptix/nerfactory into terrance/single_img_generative

docs/quickstart/installation.md

@@ -1,8 +1,35 @@
 # Installation
 
 ## Prerequisites
+::::::{tab-set}
+:::::{tab-item} Linux
 
-CUDA must be installed on the system. This library has been tested with version 11.3. You can find more information about installing CUDA [here](https://docs.nvidia.com/cuda/cuda-quick-start-guide/index.html)
+Install CUDA. This library has been tested with version 11.3. You can find CUDA download links [here](https://developer.nvidia.com/cuda-toolkit-archive) and more information about installing CUDA [here](https://docs.nvidia.com/cuda/cuda-quick-start-guide/index.html).
+
+:::::
+:::::{tab-item} Windows
+Install [Git](https://git-scm.com/downloads).
+
+Install CUDA. This library and the Windows installation process has been tested with version 11.3.
+
+::::{tab-set}
+:::{tab-item} 11.3 with Visual Studio 2019 v16.9
+
+Install Visual Studio 2019 ver 16.9. This must be done before installing CUDA. The necessary components are included in the `Desktop Development with C++` workflow (also called `C++ Build Tools` in the BuildTools edition). You can find older versions of Visual Studio 2019 [here](https://learn.microsoft.com/en-us/visualstudio/releases/2019/history#release-dates-and-build-numbers). The latest version, 16.11, may cause errors when installing tiny-cuda-nn.
+
+Install CUDA 11.3. You can find CUDA download links [here](https://developer.nvidia.com/cuda-toolkit-archive) and more information about installing CUDA [here](https://docs.nvidia.com/cuda/cuda-quick-start-guide/index.html).
+
+:::
+:::{tab-item} 11.7 with Visual Studio 2022
+
+Install Visual Studio 2022. This must be done before installing CUDA. The necessary components are included in the `Desktop Development with C++` workflow (also called `C++ Build Tools` in the BuildTools edition).
+
+Install CUDA 11.7. You can find CUDA download links [here](https://developer.nvidia.com/cuda-toolkit-archive) and more information about installing CUDA [here](https://docs.nvidia.com/cuda/cuda-quick-start-guide/index.html).
+
+:::
+::::
+:::::
+::::::
 
 ## Create environment
 
@@ -20,7 +47,8 @@ python -m pip install --upgrade pip
 (pytorch)=
 ### pytorch
 
-Either pytorch 1.12.1 (with CUDA up to 11.3) or 1.13.1 (with CUDA 11.6/7/8) can be used.
+::::{tab-set}
+:::{tab-item} Torch 1.12.1 with CUDA 11.3
 
 - To install 1.12.1 with CUDA 11.3:
 
@@ -28,6 +56,9 @@ Either pytorch 1.12.1 (with CUDA up to 11.3) or 1.13.1 (with CUDA 11.6/7/8) can
 pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 -f https://download.pytorch.org/whl/torch_stable.html
 ```
 
+:::
+:::{tab-item} Torch 1.13.1 with CUDA 11.7
+
 - To install 1.13.1 with CUDA 11.7:
 
 Note that if a pytorch version prior to 1.13 is installed, 
@@ -37,12 +68,15 @@ it should be uninstalled first to avoid upgrade issues (e.g. with functorch)
 pip uninstall torch torchvision functorch
 ```
 
-Install pytorch 1.13.1 with CUDA (this repo has been tested with CUDA 11.7.1) and [tiny-cuda-nn](https://github.com/NVlabs/tiny-cuda-nn)
+Install pytorch 1.13.1 with CUDA and [tiny-cuda-nn](https://github.com/NVlabs/tiny-cuda-nn)
 
 ```bash
 pip install torch torchvision functorch --extra-index-url https://download.pytorch.org/whl/cu117
 ```
 
+:::
+::::
+
 ### tinycudann
 
 After pytorch, install the torch bindings for tinycudann:
@@ -102,7 +136,7 @@ or be built from the repository using
 docker build --tag nerfstudio -f Dockerfile .
 ```
 ### Using an interactive container
-The docker container can be launched with an interactive terminal where nerfstudio commands can be entered as usual. Some parameters are required and some are strongly reocmmended for usage as following:
+The docker container can be launched with an interactive terminal where nerfstudio commands can be entered as usual. Some parameters are required and some are strongly recommended for usage as following:
 ```bash
 docker run --gpus all \                                         # Give the container access to nvidia GPU (required).
             -v /folder/of/your/data:/workspace/ \               # Mount a folder from the local machine into the container to be able to process them (required).
@@ -148,13 +182,15 @@ K80 | 37
 
 ## Installation FAQ
 
-- [TinyCUDA installation errors out with cuda mismatch](tiny-cuda-error)
+- [TinyCUDA installation errors out with cuda mismatch](tiny-cuda-mismatch-error)
+- [TinyCUDA installation errors out with no CUDA toolset found](tiny-cuda-integration-error)
+- [TinyCUDA installation errors out with syntax errors](tiny-cuda-syntax-error)
 - [Installation errors, File "setup.py" not found](pip-install-error)
 - [Runtime errors, "len(sources) > 0".](cuda-sources-error)
 
  <br />
 
-(tiny-cuda-error)=
+(tiny-cuda-mismatch-error)=
 
 **TinyCUDA installation errors out with cuda mismatch**
 
@@ -167,6 +203,67 @@ See [pytorch](pytorch) under Dependencies, above.
 
  <br />
 
+(tiny-cuda-integration-error)=
+
+**(Windows) TinyCUDA installation errors out with no CUDA toolset found**
+
+While installing tiny-cuda on Windows, you run into: `No CUDA toolset found.`
+
+**Solution**:
+
+Confirm that you have Visual Studio installed (CUDA 11.3 is only compatible up to 2019 ver 16.9, not 16.11 or 2022).
+
+Make sure CUDA Visual Studio integration is enabled. This should be done automatically by the CUDA installer if it is run after Visual Studio is installed. You can also manually enable integration.
+
+::::{tab-set}
+:::{tab-item} Visual Studio 2019
+
+To manually enable integration for Visual Studio 2019, copy all 4 files from
+
+```
+C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.3\extras\visual_studio_integration\MSBuildExtensions
+```
+
+to
+
+```
+C:\Program Files (x86)\Microsoft Visual Studio\2019\BuildTools\MSBuild\Microsoft\VC\v160\BuildCustomizations
+```
+
+:::
+:::{tab-item} Visual Studio 2022
+
+To manually enable integration for Visual Studio 2022, copy all 4 files from
+
+```
+C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.7\extras\visual_studio_integration\MSBuildExtensions
+```
+
+to
+
+```
+C:\Program Files\Microsoft Visual Studio\2022\[Community, Professional, Enterprise, or BuildTools]\MSBuild\Microsoft\VC\v160\BuildCustomizations
+```
+
+:::
+::::
+
+ <br />
+
+(tiny-cuda-syntax-error)=
+
+**(Windows) TinyCUDA installation errors out with syntax errors**
+
+While installing tiny-cuda on Windows, you run into: `Expected a "("`
+
+**Solution**:
+
+If you are using CUDA 11.3 with Visual Studio 2019, confirm that Visual Studio is on version 16.9 and not 16.11.
+
+If you are using a different version of CUDA and Visual Studio, try switching to CUDA 11.3 with Visual Studio 2019 ver 16.9.
+
+ <br />
+
 (pip-install-error)=
 
 **Installation errors, File "setup.py" not found**

nerfstudio/cameras/camera_optimizers.py

@@ -31,7 +31,10 @@ from typing_extensions import Literal, assert_never
 from nerfstudio.cameras.lie_groups import exp_map_SE3, exp_map_SO3xR3
 from nerfstudio.configs.base_config import InstantiateConfig
 from nerfstudio.engine.optimizers import AdamOptimizerConfig
-from nerfstudio.engine.schedulers import SchedulerConfig
+from nerfstudio.engine.schedulers import (
+    ExponentialDecaySchedulerConfig,
+    SchedulerConfig,
+)
 from nerfstudio.utils import poses as pose_utils
 
 
@@ -53,7 +56,7 @@ class CameraOptimizerConfig(InstantiateConfig):
     optimizer: AdamOptimizerConfig = AdamOptimizerConfig(lr=6e-4, eps=1e-15)
     """ADAM parameters for camera optimization."""
 
-    scheduler: SchedulerConfig = SchedulerConfig(max_steps=10000)
+    scheduler: SchedulerConfig = ExponentialDecaySchedulerConfig(max_steps=10000)
     """Learning rate scheduler for camera optimizer.."""
 
     param_group: tyro.conf.Suppress[str] = "camera_opt"

nerfstudio/configs/method_configs.py

@@ -47,7 +47,7 @@ from nerfstudio.data.dataparsers.phototourism_dataparser import (
     PhototourismDataParserConfig,
 )
 from nerfstudio.engine.optimizers import AdamOptimizerConfig, RAdamOptimizerConfig
-from nerfstudio.engine.schedulers import SchedulerConfig
+from nerfstudio.engine.schedulers import ExponentialDecaySchedulerConfig
 from nerfstudio.engine.trainer import TrainerConfig
 from nerfstudio.field_components.temporal_distortions import TemporalDistortionKind
 from nerfstudio.models.depth_nerfacto import DepthNerfactoModelConfig
@@ -274,11 +274,11 @@ method_configs["tensorf"] = TrainerConfig(
     optimizers={
         "fields": {
             "optimizer": AdamOptimizerConfig(lr=0.001),
-            "scheduler": SchedulerConfig(lr_final=0.0001, max_steps=30000),
+            "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=30000),
         },
         "encodings": {
             "optimizer": AdamOptimizerConfig(lr=0.02),
-            "scheduler": SchedulerConfig(lr_final=0.002, max_steps=30000),
+            "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.002, max_steps=30000),
         },
     },
     viewer=ViewerConfig(num_rays_per_chunk=1 << 15),
@@ -351,21 +351,22 @@ method_configs["dreamfusion"] = TrainerConfig(
             dataparser=NerfstudioDataParserConfig(),
             train_num_rays_per_batch=4096,
             eval_num_rays_per_batch=4096,
+            horizontal_rotation_warmup=4000,
         ),
         model=DreamFusionModelConfig(
             eval_num_rays_per_chunk=1 << 15,
             distortion_loss_mult=0.02,
-            interlevel_loss_mult=1.0,
+            interlevel_loss_mult=100.0,
             orientation_loss_mult=0.1,
             max_res=256,
             sphere_collider=True,
             initialize_density=False,
             random_background=True,
-            proposal_warmup=1000,
-            proposal_update_every=2,
+            proposal_warmup=500,
+            proposal_update_every=0,
             proposal_weights_anneal_max_num_iters=100,
-            start_lambertian_training=1000,
-            start_normals_training=500,
+            start_lambertian_training=500,
+            start_normals_training=1000,
             opacity_loss_mult=0.001,
             positional_prompting="discrete",
             guidance_scale=100,
@@ -373,12 +374,14 @@ method_configs["dreamfusion"] = TrainerConfig(
     ),
     optimizers={
         "proposal_networks": {
-            "optimizer": AdamOptimizerConfig(lr=5e-2, eps=1e-15),
+            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
             "scheduler": None,
         },
         "fields": {
-            "optimizer": AdamOptimizerConfig(lr=5e-3, eps=1e-15),
-            "scheduler": None,
+            "optimizer": AdamOptimizerConfig(lr=1e-3, eps=1e-15),
+            "scheduler": ExponentialDecaySchedulerConfig(
+                warmup_steps=2000, lr_final=1e-6, max_steps=20000, ramp="linear"
+            ),
         },
     },
     viewer=ViewerConfig(num_rays_per_chunk=1 << 15),

nerfstudio/data/datamanagers/dreamfusion_datamanager.py

@@ -66,7 +66,7 @@ def random_train_pose(
     radius_mean: float = 1.0,
     radius_std: float = 0.1,
     central_rotation_range: Tuple[float, float] = (0, 360),
-    vertical_rotation_range: Tuple[float, float] = (-80, 20),
+    vertical_rotation_range: Tuple[float, float] = (-90, 0),
     focal_range: Tuple[float, float] = (0.75, 1.35),
     jitter_std: float = 0.01,
     center: Tuple[float, float, float] = (0, 0, 0),
@@ -169,11 +169,14 @@ class DreamFusionDataManagerConfig(VanillaDataManagerConfig):
     """Std of radius of camera orbit"""
     focal_range: Tuple[float, float] = (0.6, 1.2)
     """Range of focal length"""
-    vertical_rotation_range: Tuple[float, float] = (-80, 30)
+    vertical_rotation_range: Tuple[float, float] = (-90, 0)
     """Range of vertical rotation"""
     jitter_std: float = 0.05
     """Std of camera direction jitter, so we don't just point the cameras towards the center every time"""
     center: Tuple[float, float, float] = (0, 0, 0)
+    """Center coordinate of the camera sphere"""
+    horizontal_rotation_warmup: int = 0
+    """How many steps until the full horizontal rotation range is used"""
 
 
 class DreamFusionDataManager(VanillaDataManager):  # pylint: disable=abstract-method
@@ -244,7 +247,8 @@ class DreamFusionDataManager(VanillaDataManager):  # pylint: disable=abstract-me
         #     ).flatten()
         #     return ray_bundle, {"initialization": False}
 
-        # TODO below
+        horizontal_range = min((step / self.config.horizontal_rotation_warmup), 1) * 180
+
         cameras, vertical_rotation, central_rotation = random_train_pose(
             self.config.train_images_per_batch,
             self.config.train_resolution,
@@ -255,6 +259,7 @@ class DreamFusionDataManager(VanillaDataManager):  # pylint: disable=abstract-me
             vertical_rotation_range=self.config.vertical_rotation_range,
             jitter_std=self.config.jitter_std,
             center=self.config.center,
+            central_rotation_range=(-horizontal_range, horizontal_range),
         )
         ray_bundle = cameras.generate_rays(torch.tensor(list(range(self.config.train_images_per_batch)))).flatten()
 

nerfstudio/engine/optimizers.py

@@ -86,8 +86,10 @@ class Optimizers:
             self.optimizers[param_group_name] = config[param_group_name]["optimizer"].setup(params=params)
             self.parameters[param_group_name] = params
             if config[param_group_name]["scheduler"]:
-                self.schedulers[param_group_name] = config[param_group_name]["scheduler"].setup(
-                    optimizer=self.optimizers[param_group_name], lr_init=lr_init
+                self.schedulers[param_group_name] = (
+                    config[param_group_name]["scheduler"]
+                    .setup()
+                    .get_scheduler(optimizer=self.optimizers[param_group_name], lr_init=lr_init)
                 )
 
     def optimizer_step(self, param_group_name: str) -> None:

nerfstudio/engine/schedulers.py

@@ -12,105 +12,98 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# pylint: disable=too-few-public-methods
+
 """Scheduler Classes"""
 
+from abc import abstractmethod
 from dataclasses import dataclass, field
-from typing import Any, Optional, Type
+from typing import Optional, Type
 
 import numpy as np
 from torch.optim import Optimizer, lr_scheduler
+from typing_extensions import Literal
 
 from nerfstudio.configs.base_config import InstantiateConfig
 
 
 @dataclass
 class SchedulerConfig(InstantiateConfig):
-    """Basic scheduler config with self-defined exponential decay schedule"""
+    """Basic scheduler config"""
 
-    _target: Type = field(default_factory=lambda: ExponentialDecaySchedule)
-    lr_final: float = 0.000005
-    max_steps: int = 1000000
-
-    # TODO: somehow make this more generic. i dont like the idea of overriding the setup function
-    # but also not sure how to go about passing things into predefined torch objects.
-    def setup(self, optimizer=None, lr_init=None, **kwargs) -> Any:
-        """Returns the instantiated object using the config."""
-        return self._target(optimizer, lr_init, self.lr_final, self.max_steps)
+    _target: Type = field(default_factory=lambda: Scheduler)
+    """target class to instantiate"""
 
 
-class ExponentialDecaySchedule(lr_scheduler.LambdaLR):
-    """Exponential learning rate decay function.
-    See https://github.com/google-research/google-research/blob/
-    fd2cea8cdd86b3ed2c640cbe5561707639e682f3/jaxnerf/nerf/utils.py#L360
-    for details.
-
-    Args:
-        optimizer: The optimizer to update.
-        lr_init: The initial learning rate.
-        lr_final: The final learning rate.
-        max_steps: The maximum number of steps.
-        lr_delay_steps: The number of steps to delay the learning rate.
-        lr_delay_mult: The multiplier for the learning rate after the delay.
-    """
+class Scheduler:
+    """Base scheduler"""
 
     config: SchedulerConfig
 
-    def __init__(self, optimizer, lr_init, lr_final, max_steps, lr_delay_steps=0, lr_delay_mult=1.0) -> None:
+    def __init__(self, config: SchedulerConfig):
+        super().__init__()
+        self.config = config
+
+    @abstractmethod
+    def get_scheduler(self, optimizer: Optimizer, lr_init: float) -> lr_scheduler._LRScheduler:
+        """Abstract method that returns a scheduler object.
+
+        Args:
+            optimizer: The optimizer to use.
+            lr_init: The initial learning rate.
+        Returns:
+            The scheduler object.
+        """
+
+
+@dataclass
+class ExponentialDecaySchedulerConfig(SchedulerConfig):
+    """Config for exponential decay scheduler with warmup"""
+
+    _target: Type = field(default_factory=lambda: ExponentialDecayScheduler)
+    """target class to instantiate"""
+    lr_pre_warmup: float = 1e-8
+    """Learning rate before warmup."""
+    lr_final: Optional[float] = None
+    """Final learning rate. If not provided, it will be set to the optimizers learning rate."""
+    warmup_steps: int = 0
+    """Number of warmup steps."""
+    max_steps: int = 100000
+    """The maximum number of steps."""
+    ramp: Literal["linear", "cosine"] = "cosine"
+    """The ramp function to use during the warmup."""
+
+
+class ExponentialDecayScheduler(Scheduler):
+    """Exponential decay scheduler with linear warmup. Scheduler first ramps up to `lr_init` in `warmup_steps`
+    steps, then exponentially decays to `lr_final` in `max_steps` steps.
+    """
+
+    config: ExponentialDecaySchedulerConfig
+
+    def get_scheduler(self, optimizer: Optimizer, lr_init: float) -> lr_scheduler._LRScheduler:
+        if self.config.lr_final is None:
+            lr_final = lr_init
+        else:
+            lr_final = self.config.lr_final
+
         def func(step):
-            if lr_delay_steps > 0:
-                delay_rate = lr_delay_mult + (1 - lr_delay_mult) * np.sin(
-                    0.5 * np.pi * np.clip(step / lr_delay_steps, 0, 1)
-                )
+            if step < self.config.warmup_steps:
+                if self.config.ramp == "cosine":
+                    lr = self.config.lr_pre_warmup + (1 - self.config.lr_pre_warmup) * np.sin(
+                        0.5 * np.pi * np.clip(step / self.config.warmup_steps, 0, 1)
+                    )
+                else:
+                    lr = (
+                        self.config.lr_pre_warmup
+                        + (lr_init - self.config.lr_pre_warmup) * step / self.config.warmup_steps
+                    )
             else:
-                delay_rate = 1.0
-            t = np.clip(step / max_steps, 0, 1)
-            log_lerp = np.exp(np.log(lr_init) * (1 - t) + np.log(lr_final) * t)
-            multiplier = (
-                log_lerp / lr_init
-            )  # divided by lr_init because the multiplier is with the initial learning rate
-            return delay_rate * multiplier
+                t = np.clip(
+                    (step - self.config.warmup_steps) / (self.config.max_steps - self.config.warmup_steps), 0, 1
+                )
+                lr = np.exp(np.log(lr_init) * (1 - t) + np.log(lr_final) * t)
+            return lr / lr_init  # divided by lr_init because the multiplier is with the initial learning rate
 
-        super().__init__(optimizer, lr_lambda=func)
-
-
-class DelayerScheduler(lr_scheduler.LambdaLR):
-    """Starts with a flat lr schedule until it reaches N epochs then applies a given scheduler"""
-
-    def __init__(
-        self,
-        optimizer: Optimizer,
-        lr_init,  # pylint: disable=unused-argument
-        lr_final,  # pylint: disable=unused-argument
-        max_steps,  # pylint: disable=unused-argument
-        delay_epochs: int = 500,
-        after_scheduler: Optional[lr_scheduler.LambdaLR] = None,
-    ) -> None:
-        def func(step):
-            if step > delay_epochs:
-                if after_scheduler is not None:
-                    multiplier = after_scheduler.lr_lambdas[0](step - delay_epochs)  # type: ignore
-                    return multiplier
-                return 1.0
-            return 0.0
-
-        super().__init__(optimizer, lr_lambda=func)
-
-
-class DelayedExponentialScheduler(DelayerScheduler):
-    """Delayer Scheduler with an Exponential Scheduler initialized afterwards."""
-
-    def __init__(
-        self,
-        optimizer: Optimizer,
-        lr_init,
-        lr_final,
-        max_steps,
-        delay_epochs: int = 200,
-    ):
-        after_scheduler = ExponentialDecaySchedule(
-            optimizer,
-            lr_init,
-            lr_final,
-            max_steps,
-        )
-        super().__init__(optimizer, lr_init, lr_final, max_steps, delay_epochs, after_scheduler=after_scheduler)
+        scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=func)
+        return scheduler

nerfstudio/generative/stable_diffusion_utils.py

@@ -106,7 +106,7 @@ class PositionalTextEmbeddings:
             text_embedding = text_embedding / 90.0
             text_embedding = (vert * text_embedding + (90 - vert) * self.top_embed) / 90.0
         else:
-            print('here')
+            print("here")
             text_embedding = self.base_embed
 
         return text_embedding

nerfstudio/model_components/ray_samplers.py

@@ -393,7 +393,6 @@ class VolumetricSampler(Sampler):
         self.occupancy_grid = occupancy_grid
         if self.scene_aabb is not None:
             self.scene_aabb = self.scene_aabb.to("cuda").flatten()
-        print(self.scene_aabb)
 
     def get_sigma_fn(self, origins, directions) -> Optional[Callable]:
         """Returns a function that returns the density of a point.

nerfstudio/models/dreamfusion.py

@@ -77,8 +77,10 @@ class DreamFusionModelConfig(ModelConfig):
     """Randomizes light source per output."""
     initialize_density: bool = True
     """Initialize density in center of scene."""
-    init_density_strength: float = 10
-    """Initial strength of center density"""
+    taper_range: Tuple[int, int] = (0, 1000)
+    """Range of step values for the density tapering"""
+    taper_strength: Tuple[float, float] = (1.0, 0.0)
+    """Strength schedule of center density"""
     sphere_collider: bool = True
     """Use spherical collider instead of box"""
     random_background: bool = True
@@ -135,7 +137,7 @@ class DreamFusionModelConfig(ModelConfig):
     """enables location based prompting"""
     interpolated_prompting: bool = False
     """enables interpolated location prompting"""
-    positional_prompting: Literal["discrete", "interpolated", "base"] = "discrete"
+    positional_prompting: Literal["discrete", "interpolated", "off"] = "discrete"
     """ how to incorporate position into prompt"""
     top_prompt: str = ", overhead view"
     """appended to prompt for overhead view"""
@@ -237,7 +239,7 @@ class DreamFusionModel(Model):
             single_jitter=self.config.use_single_jitter,
             update_sched=update_schedule,
             initial_sampler=UniformSampler(single_jitter=self.config.use_single_jitter),
-            pdf_histogram_padding=0.001,
+            pdf_histogram_padding=0.005,
         )
 
         # renderers
@@ -267,7 +269,7 @@ class DreamFusionModel(Model):
         def taper_density(
             self, training_callback_attributes: TrainingCallbackAttributes, step: int  # pylint: disable=unused-argument
         ):
-            self.density_strength -= 0.25
+            self.density_strength = np.interp(step, self.config.taper_range, self.config.taper_strength)
 
         def start_training_normals(
             self, training_callback_attributes: TrainingCallbackAttributes, step: int  # pylint: disable=unused-argument
@@ -301,8 +303,8 @@ class DreamFusionModel(Model):
         callbacks = [
             TrainingCallback(
                 where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION],
-                iters=(150, 300, 450, 600),
                 func=taper_density,
+                update_every_num_iters=1,
                 args=[self, training_callback_attributes],
             ),
             TrainingCallback(
@@ -351,12 +353,8 @@ class DreamFusionModel(Model):
 
         if self.initialize_density:
             pos = ray_samples.frustums.get_positions()
-            density_blob = (
-                self.config.init_density_strength
-                * self.density_strength
-                * torch.exp(-torch.norm(pos, dim=-1) * 10)[..., None]
-            )
-            density += density_blob
+            density_blob = self.density_strength * torch.exp(-torch.norm(pos, dim=-1) / (2 * 0.04))[..., None]
+            density = density + density_blob
 
         weights = ray_samples.get_weights(density)
         weights_list.append(weights)
@@ -425,22 +423,19 @@ class DreamFusionModel(Model):
         else:
             outputs["train_output"] = outputs["rgb"]
 
-        if self.training:
-            outputs["rendered_orientation_loss"] = orientation_loss(
-                weights.detach(), field_outputs[FieldHeadNames.NORMALS], ray_bundle.directions
-            )
+        outputs["rendered_orientation_loss"] = orientation_loss(
+            weights.detach(), field_outputs[FieldHeadNames.NORMALS], ray_bundle.directions
+        )
 
-            outputs["rendered_pred_normal_loss"] = pred_normal_loss(
-                weights.detach(),
-                field_outputs[FieldHeadNames.NORMALS].detach(),
-                field_outputs[FieldHeadNames.PRED_NORMALS],
-            )
+        outputs["rendered_pred_normal_loss"] = pred_normal_loss(
+            weights.detach(),
+            field_outputs[FieldHeadNames.NORMALS].detach(),
+            field_outputs[FieldHeadNames.PRED_NORMALS],
+        )
 
-            assert weights.shape[-1] == 1
-            if self.config.opacity_penalty:
-                outputs["opacity_loss"] = (
-                    torch.sqrt(torch.sum(weights, dim=-2) ** 2 + 0.01) * self.config.opacity_loss_mult
-                )
+        assert weights.shape[-1] == 1
+        if self.config.opacity_penalty:
+            outputs["opacity_loss"] = torch.sqrt(torch.sum(weights, dim=-2) ** 2 + 0.01) * self.config.opacity_loss_mult
 
         return outputs
 

nerfstudio/models/nerfacto.py

@@ -46,7 +46,10 @@ from nerfstudio.model_components.losses import (
     orientation_loss,
     pred_normal_loss,
 )
-from nerfstudio.model_components.ray_samplers import ProposalNetworkSampler
+from nerfstudio.model_components.ray_samplers import (
+    ProposalNetworkSampler,
+    UniformSampler,
+)
 from nerfstudio.model_components.renderers import (
     AccumulationRenderer,
     DepthRenderer,
@@ -95,6 +98,8 @@ class NerfactoModelConfig(ModelConfig):
         ]
     )
     """Arguments for the proposal density fields."""
+    proposal_initial_sampler: Literal["piecewise", "uniform"] = "piecewise"
+    """Initial sampler for the proposal network. Piecewise is preferred for unbounded scenes."""
     interlevel_loss_mult: float = 1.0
     """Proposal loss multiplier."""
     distortion_loss_mult: float = 0.002
@@ -171,12 +176,19 @@ class NerfactoModel(Model):
             1,
             self.config.proposal_update_every,
         )
+
+        # Change proposal network initial sampler if uniform
+        initial_sampler = None  # None is for piecewise as default (see ProposalNetworkSampler)
+        if self.config.proposal_initial_sampler == "uniform":
+            initial_sampler = UniformSampler(single_jitter=self.config.use_single_jitter)
+
         self.proposal_sampler = ProposalNetworkSampler(
             num_nerf_samples_per_ray=self.config.num_nerf_samples_per_ray,
             num_proposal_samples_per_ray=self.config.num_proposal_samples_per_ray,
             num_proposal_network_iterations=self.config.num_proposal_iterations,
             single_jitter=self.config.use_single_jitter,
             update_sched=update_schedule,
+            initial_sampler=initial_sampler,
         )
 
         # Collider
@@ -312,7 +324,6 @@ class NerfactoModel(Model):
     def get_image_metrics_and_images(
         self, outputs: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor]
     ) -> Tuple[Dict[str, float], Dict[str, torch.Tensor]]:
-
         image = batch["image"].to(self.device)
         rgb = outputs["rgb"]
         acc = colormaps.apply_colormap(outputs["accumulation"])

nerfstudio/models/tensorf.py

@@ -130,9 +130,13 @@ class TensoRFModel(Model):
                 "optimizer"
             ].setup(params=enc)
             if optimizers_config["encodings"]["scheduler"]:
-                training_callback_attributes.optimizers.schedulers["encodings"] = optimizers_config["encodings"][
-                    "scheduler"
-                ].setup(optimizer=training_callback_attributes.optimizers.optimizers["encodings"], lr_init=lr_init)
+                training_callback_attributes.optimizers.schedulers["encodings"] = (
+                    optimizers_config["encodings"]["scheduler"]
+                    .setup()
+                    .get_scheduler(
+                        optimizer=training_callback_attributes.optimizers.optimizers["encodings"], lr_init=lr_init
+                    )
+                )
 
         callbacks = [
             TrainingCallback(

nerfstudio/process_data/equirect_utils.py

@@ -0,0 +1,120 @@
+# Copyright 2022 The Nerfstudio Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Helper utils for processing equirectangular data."""
+
+import os
+from pathlib import Path
+from typing import Tuple
+
+import cv2
+import numpy as np
+import torch
+from equilib import Equi2Pers
+from rich.progress import (
+    BarColumn,
+    Progress,
+    TaskProgressColumn,
+    TextColumn,
+    TimeRemainingColumn,
+)
+
+from nerfstudio.utils.rich_utils import ItersPerSecColumn
+
+
+def generate_planar_projections_from_equirectangular(
+    image_dir: Path,
+    planar_image_size: Tuple[int, int],
+    samples_per_im: int,
+) -> Path:
+    """Generate planar projections from an equirectangular image.
+
+    Args:
+        image_dir: The directory containing the equirectangular image.
+        planar_image_size: The size of the planar projections [width, height].
+        samples_per_im: The number of samples to take per image.
+    returns:
+        The path to the planar projections directory.
+    """
+
+    device = torch.device("cuda")
+
+    fov = 120
+    yaw_pitch_pairs = []
+    if samples_per_im == 8:
+        fov = 120
+        for i in np.arange(-180, 180, 90):
+            yaw_pitch_pairs.append((i, 0))
+        for i in np.arange(-180, 180, 180):
+            yaw_pitch_pairs.append((i, 45))
+        for i in np.arange(-180, 180, 180):
+            yaw_pitch_pairs.append((i, -45))
+    elif samples_per_im == 14:
+        fov = 110
+        for i in np.arange(-180, 180, 60):
+            yaw_pitch_pairs.append((i, 0))
+        for i in np.arange(-180, 180, 90):
+            yaw_pitch_pairs.append((i, 45))
+        for i in np.arange(-180, 180, 90):
+            yaw_pitch_pairs.append((i, -45))
+
+    equi2pers = Equi2Pers(height=planar_image_size[1], width=planar_image_size[0], fov_x=fov, mode="bilinear")
+    frame_dir = image_dir
+    output_dir = image_dir / "planar_projections"
+    output_dir.mkdir(exist_ok=True)
+    num_ims = len(os.listdir(frame_dir))
+    progress = Progress(
+        TextColumn("[bold blue]Generating Planar Images", justify="right"),
+        BarColumn(),
+        TaskProgressColumn(show_speed=True),
+        ItersPerSecColumn(suffix="equirect frames/s"),
+        TimeRemainingColumn(elapsed_when_finished=True, compact=True),
+    )
+
+    curr_im = 0
+    with progress:
+        for i in progress.track(os.listdir(frame_dir), description="", total=num_ims):
+            if i.lower().endswith((".jpg", ".png", ".jpeg")):
+                im = np.array(cv2.imread(os.path.join(frame_dir, i)))
+                im = torch.tensor(im, dtype=torch.float32, device=device)
+                im = torch.permute(im, (2, 0, 1)) / 255.0
+                count = 0
+                for u_deg, v_deg in yaw_pitch_pairs:
+                    v_rad = torch.pi * v_deg / 180.0
+                    u_rad = torch.pi * u_deg / 180.0
+                    pers_image = equi2pers(im, rots={"roll": 0, "pitch": v_rad, "yaw": u_rad}) * 255.0
+                    pers_image = (pers_image.permute(1, 2, 0)).type(torch.uint8).to("cpu").numpy()
+                    cv2.imwrite(f"{output_dir}/{i[:-4]}_{count}.jpg", pers_image)
+                    count += 1
+            curr_im += 1
+
+    return output_dir
+
+
+def compute_resolution_from_equirect(image_dir: Path, num_images: int) -> Tuple[int, int]:
+    """Compute the resolution of the persepctive projections of equirectangular images
+       from the heuristic: num_image * res**2 = orig_height * orig_width.
+
+    Args:
+        image_dir: The directory containing the equirectangular images.
+    returns:
+        The target resolution of the perspective projections.
+    """
+
+    for i in os.listdir(image_dir):
+        if i.lower().endswith((".jpg", ".png", ".jpeg")):
+            im = np.array(cv2.imread(os.path.join(image_dir, i)))
+            res_squared = (im.shape[0] * im.shape[1]) / num_images
+            return (int(np.sqrt(res_squared)), int(np.sqrt(res_squared)))
+    raise ValueError("No images found in the directory.")

nerfstudio/process_data/process_data_utils.py

@@ -43,6 +43,7 @@ class CameraModel(Enum):
 CAMERA_MODELS = {
     "perspective": CameraModel.OPENCV,
     "fisheye": CameraModel.OPENCV_FISHEYE,
+    "equirectangular": CameraModel.OPENCV,
 }
 
 

pyproject.toml

@@ -38,6 +38,7 @@ dependencies = [
     "Pillow>=9.3.0",
     "plotly>=5.7.0",
     "protobuf<=3.20.3,!=3.20.0",
+    "pyequilib>=0.5.6",
     "pymeshlab>=2022.2.post2",
     "pyngrok>=5.1.0",
     "python-socketio>=5.7.1",

scripts/benchmarking/launch_train_blender.sh

@@ -28,7 +28,7 @@ fi
 method_opts=()
 if [ "$method_name" = "nerfacto" ]; then
     # https://github.com/nerfstudio-project/nerfstudio/issues/806#issuecomment-1284327844
-    method_opts=(--pipeline.model.background-color white --pipeline.model.near-plane 2. --pipeline.model.far-plane 6. --pipeline.datamanager.camera-optimizer.mode off --pipeline.model.use-average-appearance-embedding False)
+    method_opts=(--pipeline.model.background-color white --pipeline.model.proposal-initial-sampler uniform --pipeline.model.near-plane 2. --pipeline.model.far-plane 6. --pipeline.datamanager.camera-optimizer.mode off --pipeline.model.use-average-appearance-embedding False)
 fi
 
 shift $((OPTIND-1))

scripts/process_data.py

@@ -3,6 +3,7 @@
 
 
 import json
+import shutil
 import sys
 import zipfile
 from dataclasses import dataclass
@@ -16,6 +17,7 @@ from typing_extensions import Annotated, Literal
 
 from nerfstudio.process_data import (
     colmap_utils,
+    equirect_utils,
     hloc_utils,
     insta360_utils,
     metashape_utils,
@@ -45,7 +47,7 @@ class ProcessImages:
     """Path the data, either a video file or a directory of images."""
     output_dir: Path
     """Path to the output directory."""
-    camera_type: Literal["perspective", "fisheye"] = "perspective"
+    camera_type: Literal["perspective", "fisheye", "equirectangular"] = "perspective"
     """Camera model to use."""
     matching_method: Literal["exhaustive", "sequential", "vocab_tree"] = "vocab_tree"
     """Feature matching method to use. Vocab tree is recommended for a balance of speed and
@@ -83,6 +85,10 @@ class ProcessImages:
     """
     colmap_cmd: str = "colmap"
     """How to call the COLMAP executable."""
+    images_per_equirect: Literal[8, 14] = 8
+    """Number of samples per image to take from each equirectangular image.
+       Used only when camera-type is equirectangular.
+    """
     gpu: bool = True
     """If True, use GPU."""
     verbose: bool = False
@@ -110,6 +116,14 @@ class ProcessImages:
         image_dir = self.output_dir / "images"
         image_dir.mkdir(parents=True, exist_ok=True)
 
+        # Generate planar projections if equirectangular
+        if self.camera_type == "equirectangular":
+            pers_size = equirect_utils.compute_resolution_from_equirect(self.data, self.images_per_equirect)
+            CONSOLE.log(f"Generating {self.images_per_equirect} {pers_size} sized images per equirectangular image")
+            self.data = equirect_utils.generate_planar_projections_from_equirectangular(
+                self.data, pers_size, self.images_per_equirect
+            )
+
         summary_log = []
 
         # Copy and downscale images
@@ -208,7 +222,7 @@ class ProcessVideo:
     """Path to the output directory."""
     num_frames_target: int = 300
     """Target number of frames to use for the dataset, results may not be exact."""
-    camera_type: Literal["perspective", "fisheye"] = "perspective"
+    camera_type: Literal["perspective", "fisheye", "equirectangular"] = "perspective"
     """Camera model to use."""
     matching_method: Literal["exhaustive", "sequential", "vocab_tree"] = "vocab_tree"
     """Feature matching method to use. Vocab tree is recommended for a balance of speed and
@@ -240,6 +254,10 @@ class ProcessVideo:
     """If True, skips COLMAP and generates transforms.json if possible."""
     colmap_cmd: str = "colmap"
     """How to call the COLMAP executable."""
+    images_per_equirect: Literal[8, 14] = 8
+    """Number of samples per image to take from each equirectangular image.
+       Used only when camera-type is equirectangular.
+    """
     percent_radius_crop: float = 1.0
     """Create circle crop mask. The radius is the percent of the image diagonal."""
     percent_crop: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0)
@@ -260,9 +278,37 @@ class ProcessVideo:
 
         summary_log = []
         # Convert video to images
-        summary_log, num_extracted_frames = process_data_utils.convert_video_to_images(
-            self.data, image_dir=image_dir, num_frames_target=self.num_frames_target, verbose=self.verbose
-        )
+        if self.camera_type == "equirectangular":
+            # create temp images folder to store the equirect and perspective images
+            temp_image_dir = self.output_dir / "temp_images"
+            temp_image_dir.mkdir(parents=True, exist_ok=True)
+            summary_log, num_extracted_frames = process_data_utils.convert_video_to_images(
+                self.data, image_dir=temp_image_dir, num_frames_target=self.num_frames_target, verbose=self.verbose
+            )
+        else:
+            summary_log, num_extracted_frames = process_data_utils.convert_video_to_images(
+                self.data, image_dir=image_dir, num_frames_target=self.num_frames_target, verbose=self.verbose
+            )
+
+        # Generate planar projections if equirectangular
+        if self.camera_type == "equirectangular":
+            perspective_image_size = equirect_utils.compute_resolution_from_equirect(
+                self.output_dir / "temp_images", self.images_per_equirect
+            )
+            image_dir = equirect_utils.generate_planar_projections_from_equirectangular(
+                self.output_dir / "temp_images", perspective_image_size, self.images_per_equirect
+            )
+
+            # copy the perspective images to the image directory
+            process_data_utils.copy_images(
+                self.output_dir / "temp_images" / "planar_projections",
+                image_dir=self.output_dir / "images",
+                verbose=False,
+            )
+            image_dir = self.output_dir / "images"
+
+            # remove the temp_images folder
+            shutil.rmtree(self.output_dir / "temp_images", ignore_errors=True)
 
         # Create mask
         mask_path = process_data_utils.save_mask(