diff --git a/docs/quickstart/installation.md b/docs/quickstart/installation.md
index 5d53a4d0c..81a64b024 100644
--- a/docs/quickstart/installation.md
+++ b/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)
-(tiny-cuda-error)=
+(tiny-cuda-mismatch-error)=
**TinyCUDA installation errors out with cuda mismatch**
@@ -167,6 +203,67 @@ See [pytorch](pytorch) under Dependencies, above.
+(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
+```
+
+:::
+::::
+
+
+
+(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.
+
+
+
(pip-install-error)=
**Installation errors, File "setup.py" not found**
diff --git a/nerfstudio/cameras/camera_optimizers.py b/nerfstudio/cameras/camera_optimizers.py
index 20aaef3bc..afd002a71 100644
--- a/nerfstudio/cameras/camera_optimizers.py
+++ b/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"
diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py
index c812d781f..ba8f7e8fe 100644
--- a/nerfstudio/configs/method_configs.py
+++ b/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),
diff --git a/nerfstudio/data/datamanagers/dreamfusion_datamanager.py b/nerfstudio/data/datamanagers/dreamfusion_datamanager.py
index 153913d42..48a686df1 100644
--- a/nerfstudio/data/datamanagers/dreamfusion_datamanager.py
+++ b/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()
diff --git a/nerfstudio/engine/optimizers.py b/nerfstudio/engine/optimizers.py
index 6878eccc2..786c51115 100644
--- a/nerfstudio/engine/optimizers.py
+++ b/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:
diff --git a/nerfstudio/engine/schedulers.py b/nerfstudio/engine/schedulers.py
index c545b9124..d4554c7d8 100644
--- a/nerfstudio/engine/schedulers.py
+++ b/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
diff --git a/nerfstudio/generative/stable_diffusion_utils.py b/nerfstudio/generative/stable_diffusion_utils.py
index 2618676e9..15339fcc9 100644
--- a/nerfstudio/generative/stable_diffusion_utils.py
+++ b/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
diff --git a/nerfstudio/model_components/ray_samplers.py b/nerfstudio/model_components/ray_samplers.py
index 67201bb18..63a01f4fa 100644
--- a/nerfstudio/model_components/ray_samplers.py
+++ b/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.
diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py
index 4e48a5661..4dd4b6b1f 100644
--- a/nerfstudio/models/dreamfusion.py
+++ b/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
diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py
index 4972dd5be..bf92968da 100644
--- a/nerfstudio/models/nerfacto.py
+++ b/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"])
diff --git a/nerfstudio/models/tensorf.py b/nerfstudio/models/tensorf.py
index d5af1579d..48310fde2 100644
--- a/nerfstudio/models/tensorf.py
+++ b/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(
diff --git a/nerfstudio/process_data/equirect_utils.py b/nerfstudio/process_data/equirect_utils.py
new file mode 100644
index 000000000..596af58fb
--- /dev/null
+++ b/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.")
diff --git a/nerfstudio/process_data/process_data_utils.py b/nerfstudio/process_data/process_data_utils.py
index 4919edd93..2d78469eb 100644
--- a/nerfstudio/process_data/process_data_utils.py
+++ b/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,
}
diff --git a/pyproject.toml b/pyproject.toml
index f1154fcd1..0ce681905 100644
--- a/pyproject.toml
+++ b/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",
diff --git a/scripts/benchmarking/launch_train_blender.sh b/scripts/benchmarking/launch_train_blender.sh
index 298e21749..8f7148c03 100644
--- a/scripts/benchmarking/launch_train_blender.sh
+++ b/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))
diff --git a/scripts/process_data.py b/scripts/process_data.py
index 752e1d15c..7b6bb7922 100644
--- a/scripts/process_data.py
+++ b/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(