mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-12 19:04:03 +08:00
viewer stuff
This commit is contained in:
Justin committed
1 parent
3f73e78bfd
commit
23cebc83a1
4 files changed
+47
-31
No files matched your search
@@ -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,
|
||||
),
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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]:
|
||||
|
||||
@@ -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."""
|
||||
|
||||
Reference in new issue
Block a user