From 37e763f2de52b27531a34bf22f7df4d7c426dbb1 Mon Sep 17 00:00:00 2001 From: maturk Date: Fri, 30 Jun 2023 14:09:57 +0300 Subject: [PATCH] optional pose noise in datamanager --- .../data/datamanagers/parallel_datamanager.py | 28 +++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/nerfstudio/data/datamanagers/parallel_datamanager.py b/nerfstudio/data/datamanagers/parallel_datamanager.py index f8060fd79..8261d3669 100644 --- a/nerfstudio/data/datamanagers/parallel_datamanager.py +++ b/nerfstudio/data/datamanagers/parallel_datamanager.py @@ -20,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 @@ -33,6 +34,7 @@ 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 @@ -75,6 +77,12 @@ class ParallelDataManagerConfig(DataManagerConfig): 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): @@ -163,6 +171,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) @@ -228,3 +239,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)