merge

nerfstudio/data/datamanagers/parallel_datamanager.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import concurrent.futures
 import functools
 import time
+import queue
 from abc import abstractmethod
 from collections import defaultdict
 from dataclasses import dataclass, field
@@ -19,6 +20,7 @@ from typing_extensions import TypeVar
 
 from nerfstudio.cameras.cameras import CameraType
 from nerfstudio.cameras.rays import RayBundle
+from nerfstudio.cameras.lie_groups import exp_map_SE3
 from nerfstudio.configs.base_config import InstantiateConfig
 from nerfstudio.configs.dataparser_configs import AnnotatedDataParserUnion
 from nerfstudio.data.datamanagers.base_datamanager import DataManager, DataManagerConfig, TDataset, variable_res_collate
@@ -32,11 +34,12 @@ from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttrib
 from nerfstudio.model_components.ray_generators import RayGenerator
 from nerfstudio.utils.misc import IterableWrapper
 from nerfstudio.utils.rich_utils import CONSOLE
+from nerfstudio.utils import poses as pose_utils
 
 
 @dataclass
 class ParallelDataManagerConfig(DataManagerConfig):
-    """A basic data manager"""
+    """Config for a `ParallelDataManager` which reads data in multiple processes"""
 
     _target: Type = field(default_factory=lambda: ParallelDataManager)
     """Target class to instantiate."""
@@ -68,6 +71,18 @@ class ParallelDataManagerConfig(DataManagerConfig):
     """Size of patch to sample from. If >1, patch-based sampling will be used."""
     n_procs: int = 1
     """Number of processes to use for data loading."""
+    queue_size: int = 2
+    """Size of shared data queue containing generated ray bundles and batches. 
+    If queue_size <= 0, the queue size is infinite."""
+    max_thread_workers: Optional[int] = 1
+    """Maximum number of threads to use in thread pool executor. If None, automatically
+    set to harware cpu_count + 4."""
+    position_noise_std: float = 0.0
+    """Noise to add to initial camera pose positions. Useful for debugging."""
+    orientation_noise_std: float = 0.0
+    """Noise to add to initial camera pose orientations. Useful for debugging."""
+    non_trainable_camera_indices: Optional[Tuple[int, ...]] = None
+    """List of non trainable camera indices"""
 
 
 class DataProc(mp.Process):
@@ -108,8 +123,11 @@ class DataProc(mp.Process):
                 try:
                     self.out_queue.put_nowait((ray_bundle, batch))
                     break
-                except:
-                    time.sleep(0.001)
+                except queue.Full:
+                    time.sleep(0.0001)
+                except Exception:
+                    CONSOLE.print_exception()
+                    CONSOLE.print("[bold red]Error: Error occured in parallel datamanager queue.")
 
     def cache_images(self):
         # caches all the input images in a NxHxWx3 tensor
@@ -152,6 +170,9 @@ class ParallelDataManager(DataManager, Generic[TDataset]):
             self.dataparser.downscale_factor = 1  # Avoid opening images
         self.includes_time = self.dataparser.includes_time
         self.train_dataparser_outputs: DataparserOutputs = self.dataparser.get_dataparser_outputs(split="train")
+        self.non_trainable_camera_indices = self.config.non_trainable_camera_indices
+        if self.config.position_noise_std != 0.0 or self.config.orientation_noise_std != 0.0:
+            self.apply_pose_noise(self.non_trainable_camera_indices)
         self.train_dataset = self.create_train_dataset()
         # self.eval_dataset = self.create_eval_dataset()
         self.pix_sampler = self._get_pixel_sampler(self.train_dataset, self.config.train_num_rays_per_batch)
@@ -159,7 +180,7 @@ class ParallelDataManager(DataManager, Generic[TDataset]):
         mp.set_start_method("spawn")
         # Manager().queue() can be faster
         # https://stackoverflow.com/questions/43439194/python-multiprocessing-queue-vs-multiprocessing-manager-queue/45236748#45236748
-        self.data_q = mp.Manager().Queue(maxsize=3)
+        self.data_q = mp.Manager().Queue(maxsize=self.config.queue_size)
         self.data_procs = [
             DataProc(
                 self.data_q,
@@ -174,7 +195,7 @@ class ParallelDataManager(DataManager, Generic[TDataset]):
             proc.start()
         super().__init__()
         # Prime the executor with the first batch
-        self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
+        self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=self.config.max_thread_workers)
         self.batch_fut = self.executor.submit(self.data_q.get)
 
     def create_train_dataset(self) -> TDataset:
@@ -222,3 +243,20 @@ class ParallelDataManager(DataManager, Generic[TDataset]):
             A list of dictionaries containing the data manager's param groups.
         """
         return {}
+
+    def apply_pose_noise(self, non_trainable_camera_indices: Optional[List[int]]):
+        """Apply noise to training camera poses"""
+        assert self.config.position_noise_std >= 0.0 and self.config.orientation_noise_std >= 0.0
+        camera_to_worlds = self.train_dataparser_outputs.cameras.camera_to_worlds.to(self.device)
+        num_cameras = len(self.train_dataparser_outputs.cameras)
+        std_vector = torch.tensor(
+            [self.config.position_noise_std] * 3 + [self.config.orientation_noise_std] * 3, device=self.device
+        )
+        pose_noise = exp_map_SE3(torch.normal(torch.zeros((num_cameras, 6), device=self.device), std_vector))
+        if non_trainable_camera_indices is not None:
+            pose_noise[torch.tensor(non_trainable_camera_indices).long()] = torch.eye(4, device=pose_noise.device)[
+                :3, :4
+            ]
+        self.train_dataparser_outputs.cameras.camera_to_worlds = functools.reduce(
+            pose_utils.multiply, [camera_to_worlds, pose_noise]
+        ).to(self.train_dataparser_outputs.cameras.camera_to_worlds.device)

nerfstudio/fields/nerfacto_field.py

@@ -38,6 +38,7 @@ from nerfstudio.field_components.field_heads import (
 from nerfstudio.field_components.mlp import MLP
 from nerfstudio.field_components.spatial_distortions import SpatialDistortion
 from nerfstudio.fields.base_field import Field, shift_directions_for_tcnn
+from nerfstudio.utils.writer import GLOBAL_BUFFER
 
 
 class NerfactoField(Field):
@@ -64,6 +65,8 @@ class NerfactoField(Field):
         use_pred_normals: whether to use predicted normals
         use_average_appearance_embedding: whether to use average appearance embedding or zeros for inference
         spatial_distortion: spatial distortion to apply to the scene
+        bundle_adjust: whether to use bundle adjustment
+        coarse_to_fine_iters: iterations (percentage of max iters) at which bundle adjustment is active
     """
 
     aabb: Tensor
@@ -95,6 +98,8 @@ class NerfactoField(Field):
         use_average_appearance_embedding: bool = False,
         spatial_distortion: Optional[SpatialDistortion] = None,
         implementation: Literal["tcnn", "torch"] = "tcnn",
+        bundle_adjust: bool = False,
+        coarse_to_fine_iters: Optional[Tuple[float, float]] = None,
     ) -> None:
         super().__init__()
 
@@ -116,6 +121,8 @@ class NerfactoField(Field):
         self.pass_semantic_gradients = pass_semantic_gradients
         self.base_res = base_res
         self.freq_warmup = freq_warmup
+        self.bundle_adjust = bundle_adjust
+        self.coarse_to_fine_iters = coarse_to_fine_iters
         self.step = 0
 
         self.direction_encoding = SHEncoding(
@@ -207,9 +214,24 @@ class NerfactoField(Field):
     def get_freq_mask(self, B, N, D, device):
         import numpy as np
 
-        ts = np.linspace(self.step * N, self.step, N)
-        mask_vals = np.interp(ts, [0, self.freq_warmup], [0, 1])
-        m = torch.from_numpy(mask_vals)[None, ..., None].to(device)
+        if self.bundle_adjust:
+            # set masks for hash grid levels
+            self.max_iters: int = GLOBAL_BUFFER.get("max_iter")
+            assert self.coarse_to_fine_iters is not None
+            start, end = self.coarse_to_fine_iters
+            assert (
+                start >= 0 and end >= 0
+            ), f"start and end iterations for bundle adjustment have to be positive, got start = {start} and end = {end}"
+            L = N
+            # From https://arxiv.org/pdf/2104.06405.pdf equation 14
+            alpha = ((self.step / self.max_iters) - start) / (end - start) * L
+            k = torch.arange(L, dtype=torch.float32, device=device)
+            mask_vals = (1 - (alpha - k).clamp_(min=0, max=1).mul_(np.pi).cos_()) / 2
+        else:
+            ts = np.linspace(self.step * N, self.step, N)
+            mask_vals = np.interp(ts, [0, self.freq_warmup], [0, 1])
+            mask_vals = torch.from_numpy(mask_vals)
+        m = mask_vals[None, ..., None].to(device)
         return m.repeat((B, 1, D))
 
     def mask_freqs(self, hashgrid_outputs):

nerfstudio/models/nerfacto.py

@@ -20,7 +20,7 @@ from __future__ import annotations
 
 from contextlib import nullcontext
 from dataclasses import dataclass, field
-from typing import Dict, List, Literal, Tuple, Type
+from typing import Dict, List, Literal, Tuple, Type, Optional
 
 import numpy as np
 import torch
@@ -130,6 +130,11 @@ class NerfactoModelConfig(ModelConfig):
     """Dimension of the appearance embedding."""
     camera_optimizer: CameraOptimizerConfig = CameraOptimizerConfig(mode="SO3xR3")
     """Config of the camera optimizer to use"""
+    bundle_adjust: bool = False
+    """Whether to bundle adjust (BARF)"""
+    coarse_to_fine_iters: Optional[Tuple[float, float]] = (0.0, 0.1)
+    """Iterations (as a percentage of total iterations) at which coarse to fine hash grid optimization starts and ends.
+    Linear interpolation between (start, end) and full activation of hash grid from end onwards."""
 
 
 class NerfactoModel(Model):
@@ -165,6 +170,8 @@ class NerfactoModel(Model):
             use_average_appearance_embedding=self.config.use_average_appearance_embedding,
             appearance_embedding_dim=self.config.appearance_embed_dim,
             implementation=self.config.implementation,
+            bundle_adjust=self.config.bundle_adjust,
+            coarse_to_fine_iters=self.config.coarse_to_fine_iters,
         )
 
         self.camera_optimizer: CameraOptimizer = self.config.camera_optimizer.setup(