diff --git a/scripts/viewer/run_viewer.py b/scripts/viewer/run_viewer.py index 8bbff3ee3..1c38cf10f 100644 --- a/scripts/viewer/run_viewer.py +++ b/scripts/viewer/run_viewer.py @@ -54,7 +54,52 @@ class RunViewer: config.viewer = self.viewer.as_viewer_config() config.viewer.num_rays_per_chunk = num_rays_per_chunk - _start_viewer(config, pipeline, step) + self._start_viewer(config, pipeline) + + def _start_viewer(self, config: TrainerConfig, pipeline: Pipeline): + base_dir = config.get_base_dir() + viewer_log_path = base_dir / config.viewer.relative_log_filename + + datapath = pipeline.datamanager.get_datapath() + if datapath is None: + datapath = base_dir + + viewer_state, banner_messages = viewer_utils.setup_viewer( + config.viewer, + log_filename=viewer_log_path, + datapath=datapath, + ) + + # We don't need logging, but writer.GLOBAL_BUFFER needs to be populated + config.logging.local_writer.enable = False + writer.setup_local_writer(config.logging, max_iter=config.max_num_iterations, banner_messages=banner_messages) + + assert viewer_state and pipeline.datamanager.train_dataset + viewer_state.init_scene( + dataset=pipeline.datamanager.train_dataset, + start_train=False, + ) + while True: + viewer_state.vis["renderingState/isTraining"].write(False) + self._update_viewer_state(viewer_state, pipeline) + + def _update_viewer_state(self, viewer_state: viewer_utils.ViewerState, pipeline: Pipeline): + """Updates the viewer state by rendering out scene with current pipeline + Returns the time taken to render scene. + + """ + # NOTE: step must be > 0 otherwise the rendering would not happen + step = 1 + num_rays_per_batch = pipeline.datamanager.get_train_rays_per_batch() + with TimeWriter(writer, EventName.ITER_VIS_TIME) as _: + try: + viewer_state.update_scene(self, step, pipeline.model, num_rays_per_batch) + except RuntimeError: + time.sleep(0.03) # sleep to allow buffer to reset + assert viewer_state.vis is not None + viewer_state.vis["renderingState/log_errors"].write( + "Error: GPU out of memory. Reduce resolution to prevent viewer from crashing." + ) def save_checkpoint(self, *args, **kwargs): """