diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 801ef7ca3..2db6c4686 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -26,27 +26,36 @@ import tyro from nerfstudio.cameras.camera_optimizers import CameraOptimizerConfig from nerfstudio.configs.base_config import ViewerConfig from nerfstudio.configs.external_methods import get_external_methods -from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManager, VanillaDataManagerConfig -from nerfstudio.data.datamanagers.parallel_datamanager import ParallelDataManagerConfig -from nerfstudio.data.datamanagers.random_cameras_datamanager import RandomCamerasDataManagerConfig -from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig +from nerfstudio.data.datamanagers.base_datamanager import ( + VanillaDataManager, VanillaDataManagerConfig) +from nerfstudio.data.datamanagers.parallel_datamanager import \ + ParallelDataManagerConfig +from nerfstudio.data.datamanagers.random_cameras_datamanager import \ + RandomCamerasDataManagerConfig +from nerfstudio.data.dataparsers.blender_dataparser import \ + BlenderDataParserConfig from nerfstudio.data.dataparsers.dnerf_dataparser import DNeRFDataParserConfig -from nerfstudio.data.dataparsers.instant_ngp_dataparser import InstantNGPDataParserConfig -from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataParserConfig -from nerfstudio.data.dataparsers.phototourism_dataparser import PhototourismDataParserConfig -from nerfstudio.data.dataparsers.sdfstudio_dataparser import SDFStudioDataParserConfig -from nerfstudio.data.dataparsers.sitcoms3d_dataparser import Sitcoms3DDataParserConfig +from nerfstudio.data.dataparsers.instant_ngp_dataparser import \ + InstantNGPDataParserConfig +from nerfstudio.data.dataparsers.nerfstudio_dataparser import \ + NerfstudioDataParserConfig +from nerfstudio.data.dataparsers.phototourism_dataparser import \ + PhototourismDataParserConfig +from nerfstudio.data.dataparsers.sdfstudio_dataparser import \ + SDFStudioDataParserConfig +from nerfstudio.data.dataparsers.sitcoms3d_dataparser import \ + Sitcoms3DDataParserConfig from nerfstudio.data.datasets.depth_dataset import DepthDataset from nerfstudio.data.datasets.sdf_dataset import SDFDataset from nerfstudio.data.datasets.semantic_dataset import SemanticDataset -from nerfstudio.engine.optimizers import AdamOptimizerConfig, RAdamOptimizerConfig -from nerfstudio.engine.schedulers import ( - CosineDecaySchedulerConfig, - ExponentialDecaySchedulerConfig, - MultiStepSchedulerConfig, -) +from nerfstudio.engine.optimizers import (AdamOptimizerConfig, + RAdamOptimizerConfig) +from nerfstudio.engine.schedulers import (CosineDecaySchedulerConfig, + ExponentialDecaySchedulerConfig, + MultiStepSchedulerConfig) from nerfstudio.engine.trainer import TrainerConfig -from nerfstudio.field_components.temporal_distortions import TemporalDistortionKind +from nerfstudio.field_components.temporal_distortions import \ + TemporalDistortionKind from nerfstudio.fields.sdf_field import SDFFieldConfig from nerfstudio.models.depth_nerfacto import DepthNerfactoModelConfig from nerfstudio.models.generfacto import GenerfactoModelConfig @@ -170,7 +179,7 @@ method_configs["nerfacto-huge"] = TrainerConfig( pipeline=VanillaPipelineConfig( # datamanager=VanillaDataManagerConfig( datamanager=ParallelDataManagerConfig( - dataparser=NerfstudioDataParserConfig(), + dataparser=NerfstudioDataParserConfig(train_split_fraction=.95), train_num_rays_per_batch=16384, eval_num_rays_per_batch=4096, ), diff --git a/nerfstudio/data/datamanagers/parallel_datamanager.py b/nerfstudio/data/datamanagers/parallel_datamanager.py index 57fdac1ca..60dc3a2ea 100644 --- a/nerfstudio/data/datamanagers/parallel_datamanager.py +++ b/nerfstudio/data/datamanagers/parallel_datamanager.py @@ -140,7 +140,6 @@ class ParallelDataManager(DataManager, Generic[TDataset]): self.device = device self.world_size = world_size self.local_rank = local_rank - # self.sampler = None self.test_mode = test_mode self.test_split = "test" if test_mode in ["test", "inference"] else "val" self.dataparser_config = self.config.dataparser @@ -160,9 +159,15 @@ 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=2) + self.data_q = mp.Manager().Queue(maxsize=3) self.data_procs = [ - DataProc(self.data_q, self.config, self.train_dataparser_outputs, self.train_dataset, self.pix_sampler) + DataProc( + self.data_q, + self.config, + self.train_dataparser_outputs, + self.train_dataset, + self.pix_sampler, + ) for i in range(self.config.n_procs) ] for proc in self.data_procs: @@ -195,7 +200,6 @@ class ParallelDataManager(DataManager, Generic[TDataset]): def next_train(self, step: int) -> Tuple[RayBundle, Dict]: self.train_count += 1 - # Fetch the next batch in an executor to parallelize the queue get() operation # with the train step bundle, batch = self.batch_fut.result() diff --git a/nerfstudio/viewer_beta/render_state_machine.py b/nerfstudio/viewer_beta/render_state_machine.py index bfbf9570a..a871a1e65 100644 --- a/nerfstudio/viewer_beta/render_state_machine.py +++ b/nerfstudio/viewer_beta/render_state_machine.py @@ -219,6 +219,7 @@ class RenderStateMachine(threading.Thread): image=outputs[self.viewer.control_panel.output_render], colormap_options=self.viewer.control_panel.colormap_options, ) + depth = outputs["depth"] if self.viewer.control_panel.split: split_output_render = self.viewer.control_panel.split_output_render @@ -237,11 +238,14 @@ class RenderStateMachine(threading.Thread): selected_output[:, split_index] = torch.tensor([0.133, 0.157, 0.192], device=selected_output.device) selected_output = (selected_output * 255).type(torch.uint8) + depth = (torch.clip(depth,0,1) * 2**16-1).squeeze().cpu().numpy().astype(np.uint16) + # depth = torch.concatenate([depth, depth, depth], dim=-1) self.viewer.viser_server.set_background_image( selected_output.cpu().numpy(), format=self.viewer.config.image_format, jpeg_quality=self.viewer.config.jpeg_quality, + depthimg=depth, ) def _calculate_image_res(self, aspect_ratio: float) -> Tuple[int, int]: diff --git a/nerfstudio/viewer_beta/viewer.py b/nerfstudio/viewer_beta/viewer.py index 1e51fc8e9..672d20db9 100644 --- a/nerfstudio/viewer_beta/viewer.py +++ b/nerfstudio/viewer_beta/viewer.py @@ -132,16 +132,15 @@ class Viewer: camera_optimizer = self.trainer.pipeline.model.camera_optimizer with torch.no_grad(): c2ws_delta = camera_optimizer(torch.tensor(idxs, device=camera_optimizer.device)).cpu().numpy() - with self.viser_server.atomic(): - for idx in idxs: - # both are numpy arrays - c2w_orig = self.original_c2w[idx] - c2w_delta = c2ws_delta[idx, ...] - c2w = c2w_orig @ np.concatenate((c2w_delta, np.array([[0, 0, 0, 1]])), axis=0) - R = vtf.SO3.from_matrix(c2w[:3, :3]) - R = R @ vtf.SO3.from_x_radians(np.pi) - self.camera_handles[idx].position = c2w[:3, 3] * VISER_NERFSTUDIO_SCALE_RATIO - self.camera_handles[idx].wxyz = R.wxyz + for idx in idxs: + # both are numpy arrays + c2w_orig = self.original_c2w[idx] + c2w_delta = c2ws_delta[idx, ...] + c2w = c2w_orig @ np.concatenate((c2w_delta, np.array([[0, 0, 0, 1]])), axis=0) + R = vtf.SO3.from_matrix(c2w[:3, :3]) + R = R @ vtf.SO3.from_x_radians(np.pi) + self.camera_handles[idx].position = c2w[:3, 3] * VISER_NERFSTUDIO_SCALE_RATIO + self.camera_handles[idx].wxyz = R.wxyz def _interrupt_render(self, _) -> None: """Interrupt current render."""