diff --git a/docs/tutorials/quickstart.md b/docs/tutorials/quickstart.md index 0e71ab9a3..e8ebf885d 100644 --- a/docs/tutorials/quickstart.md +++ b/docs/tutorials/quickstart.md @@ -103,7 +103,7 @@ To evaluate the trained NeRF, we provide an evaluation script that allows you to python scripts/run_eval.py render-trajectory --load-config=outputs/blender_lego/instant_ngp/2022-07-07_230905/config.yml --traj=spiral --output-path=output.mp4 ``` -Please note, this quickstart allows you to preform everything in a headless manner. We also provide a web-based visualizer that allows you to easily monitor training or render out trajectories. See our [viewer docs](../tutorials/viewer/viewer_quickstart.md) for more. +Please note, this quickstart allows you to preform everything in a headless manner. We also provide a web-based viewer that allows you to easily monitor training or render out trajectories. See our [viewer docs](../tutorials/viewer/viewer_quickstart.md) for more. # FAQ diff --git a/nerfactory/configs/base.py b/nerfactory/configs/base.py index 842ea0976..ce6f394c2 100644 --- a/nerfactory/configs/base.py +++ b/nerfactory/configs/base.py @@ -19,11 +19,10 @@ from __future__ import annotations from dataclasses import dataclass from datetime import datetime from pathlib import Path -from typing import Any, Dict, List, Literal, Optional, Tuple, Type +from typing import Any, Dict, List, Literal, Optional, Tuple, Type, Union import dcargs import torch -from rich import console from nerfactory.configs.utils import to_immutable_dict @@ -46,7 +45,6 @@ from nerfactory.optimizers.schedulers import ExponentialDecaySchedule from nerfactory.pipelines.base import Pipeline from nerfactory.utils import writer -CONSOLE = console.Console() # Pretty printing class class PrintableConfig: # pylint: disable=too-few-public-methods @@ -101,10 +99,6 @@ class TensorboardWriterConfig(InstantiateConfig): _target: Type = writer.TensorboardWriter """target class to instantiate""" - enable: bool = False - """if True enables tensorboard logging, else disables""" - relative_log_dir: Path = Path("./") - """relative path to save all tensorboard events""" log_dir: dcargs.conf.Fixed[Path] = dcargs.MISSING """auto-populated absolute path to saved tensorboard events""" @@ -115,10 +109,6 @@ class WandbWriterConfig(InstantiateConfig): _target: Type = writer.WandbWriter """target class to instantiate""" - enable: bool = False - """if True enables wandb logging, else disables""" - relative_log_dir: Path = Path("./") - """relative path to save all wandb events""" log_dir: dcargs.conf.Fixed[Path] = dcargs.MISSING """auto-populated absolute path to saved wandb events""" @@ -141,10 +131,6 @@ class LocalWriterConfig(InstantiateConfig): """specifies which stats will be logged/printed to terminal""" max_log_size: int = 10 """maximum number of rows to print before wrapping. if 0, will print everything.""" - relative_log_dir: Path = Path("./") - """relative local path to save all events""" - log_dir: dcargs.conf.Fixed[Path] = dcargs.MISSING - """auto-populated absolute local path to saved events""" def setup(self, banner_messages: Optional[List[str]] = None, **kwargs) -> Any: """Instantiate local writer @@ -159,17 +145,19 @@ class LocalWriterConfig(InstantiateConfig): class LoggingConfig(PrintableConfig): """Configuration of loggers and profilers""" + relative_log_dir: Path = Path("./") + """relative path to save all logged events""" steps_per_log: int = 10 """number of steps between logging stats""" max_buffer_size: int = 20 """maximum history size to keep for computing running averages of stats. e.g. if 20, averages will be computed over past 20 occurances.""" - writer: Tuple[Any, ...] = ( - TensorboardWriterConfig(enable=False), - WandbWriterConfig(enable=True), - LocalWriterConfig(enable=True), - ) - """list of all supported writers. Can turn on/off writers by specifying enable.""" + event_writer: Literal["tb", "wandb", "none"] = "wandb" + """specify which writer to use (tensorboard or wandb)""" + event_writer_config: dcargs.conf.Fixed[Optional[Union[TensorboardWriterConfig, WandbWriterConfig]]] = dcargs.MISSING + """dynamically set the writer config based on the writer""" + local_writer: LocalWriterConfig = LocalWriterConfig(enable=True) + """if provided, will print stats locally. if None, will disable printing""" enable_profiler: bool = True """whether to enable profiling code; prints speed of functions at the end of a program. profiler logs run times of functions and prints at end of training""" @@ -487,12 +475,12 @@ class PipelineConfig(InstantiateConfig): class ViewerConfig(PrintableConfig): """Configuration for viewer instantiation""" + enable: bool = False + """whether to enable viewer""" relative_log_filename: str = "viewer_log_filename.txt" """Filename to use for the log file.""" log_filename: dcargs.conf.Fixed[Path] = dcargs.MISSING """Absolute path to the log file. Set automatically.""" - enable: bool = False - """whether to enable viewer""" start_train: bool = True """whether to immediately start training upon loading viewer if False, will just visualize dataset but you can toggle training in viewer""" @@ -503,7 +491,7 @@ class ViewerConfig(PrintableConfig): websocket_port: int = 7007 """the default websocket port to connect to""" num_rays_per_chunk: int = 32768 - """number of rays per chunk to render with visualizer""" + """number of rays per chunk to render with viewer""" # Optimizer related configs @@ -547,6 +535,7 @@ class Config(PrintableConfig): """Experiment base directory. Set automatically.""" machine: MachineConfig = MachineConfig() logging: LoggingConfig = LoggingConfig() + viewer: ViewerConfig = ViewerConfig() trainer: TrainerConfig = TrainerConfig() pipeline: PipelineConfig = PipelineConfig() optimizers: Dict[str, Any] = to_immutable_dict( @@ -557,37 +546,35 @@ class Config(PrintableConfig): } } ) - viewer: ViewerConfig = ViewerConfig() def __post_init__(self): """Make paths more specific using the current timestamp.""" - self.set_timestamp() + self.populate_dynamic_fields() - # if the viewer is enabled, disable tensorboard and wandb logging - # TODO(ethan): this is gross for now and will require a config refactor - if self.viewer.enable: - string = "Disabling tensorboard and wandb logging since viewer is enabled." - CONSOLE.print(f"[bold red]{string}") - self.logging.writer[0].enable = False - self.logging.writer[1].enable = False - # also disable eval steps - string = "Disabling eval iterations since viewer is enabled." - CONSOLE.print(f"[bold red]{string}") - self.trainer.steps_per_eval_batch = self.trainer.max_num_iterations - self.trainer.steps_per_eval_image = self.trainer.max_num_iterations - self.trainer.steps_per_eval_all_images = self.trainer.max_num_iterations - - def set_timestamp(self, timestamp: Optional[str] = None) -> None: + def populate_dynamic_fields(self, timestamp: Optional[str] = None) -> None: """Make paths in our config more specific using a timestamp. Args: timestamp: Timestamp to use, as a string. If None, defaults to the current time in the form YYYY-MM-DD_HHMMMSS. """ + + # set the timestamp of the model logging/writer loggign paths if timestamp is None: timestamp = datetime.now().strftime("%Y-%m-%d_%H%M%S") self.base_dir = Path(f"outputs/{self.experiment_name}/{self.method_name}/{timestamp}") self.trainer.model_dir = self.base_dir / self.trainer.relative_model_dir - for curr_writer in self.logging.writer: - curr_writer.log_dir = self.base_dir / curr_writer.relative_log_dir + if self.logging.event_writer == "tb": + self.logging.event_writer_config = TensorboardWriterConfig( + log_dir=self.base_dir / self.logging.relative_log_dir + ) + elif self.logging.event_writer == "wandb": + self.logging.event_writer_config = WandbWriterConfig(log_dir=self.base_dir / self.logging.relative_log_dir) + else: + self.logging.event_writer_config = None self.viewer.log_filename = self.base_dir / self.viewer.relative_log_filename + # disable test if viewer is enabled (for speed purposes) + if self.viewer.enable: + self.trainer.steps_per_eval_batch = self.trainer.max_num_iterations + self.trainer.steps_per_eval_image = self.trainer.max_num_iterations + self.trainer.steps_per_eval_all_images = self.trainer.max_num_iterations diff --git a/nerfactory/configs/base_configs.py b/nerfactory/configs/base_configs.py index f3e0de891..255ae504e 100644 --- a/nerfactory/configs/base_configs.py +++ b/nerfactory/configs/base_configs.py @@ -180,7 +180,7 @@ def _make_base_config_subcommand_type() -> Type[Config]: base_configs_placeholder_timestamp = {} for name, config in base_configs.items(): base_configs_placeholder_timestamp[name] = copy.deepcopy(config) - base_configs_placeholder_timestamp[name].set_timestamp("{timestamp}") + base_configs_placeholder_timestamp[name].populate_dynamic_fields("{timestamp}") return dcargs.extras.subcommand_type_from_defaults(base_configs_placeholder_timestamp) diff --git a/nerfactory/engine/trainer.py b/nerfactory/engine/trainer.py index e33fc36e1..e00f77ff3 100644 --- a/nerfactory/engine/trainer.py +++ b/nerfactory/engine/trainer.py @@ -26,6 +26,7 @@ from pathlib import Path from typing import Any, Dict, List import torch +from rich import console from torch.cuda.amp.grad_scaler import GradScaler from nerfactory.configs import base as cfg @@ -37,12 +38,13 @@ from nerfactory.utils.callbacks import ( TrainingCallbackAttributes, TrainingCallbackLocation, ) -from nerfactory.utils.decorators import check_main_thread +from nerfactory.utils.decorators import check_main_thread, check_viewer_enabled from nerfactory.utils.misc import step_check from nerfactory.utils.writer import EventName, TimeWriter from nerfactory.viewer.server import viewer_utils logging.getLogger("PIL").setLevel(logging.WARNING) +CONSOLE = console.Console() def train_loop(local_rank: int, world_size: int, config: cfg.Config) -> Any: @@ -81,21 +83,19 @@ class Trainer: self.pipeline: Pipeline self.optimizers: Optimizers self.start_step = 0 - # visualizer variable - banner_messages = None - self.visualizer_state = viewer_utils.VisualizerState(config.viewer, config_base_dir=self.config.base_dir) - if config.viewer.enable: - banner_messages = [f"Viewer at: {self.visualizer_state.viewer_url}"] - self.grad_scaler = GradScaler(enabled=self.mixed_precision) # training callbacks self.callbacks: List[TrainingCallback] - # logging variables - writer.setup_event_writers( + # optimizers + self.grad_scaler = GradScaler(enabled=self.mixed_precision) + # logging/viewer variables + self.viewer_state, banner_messages = viewer_utils.setup_viewer(config.viewer) + writer.setup_event_writer(config.logging) + writer.setup_local_writer( config.logging, max_iter=config.trainer.max_num_iterations, banner_messages=banner_messages ) profiler.setup_profiler(config.logging) - writer.put_config(name="config", config_dict=dataclasses.asdict(config), step=0) + self._check_viewer_warnings() def setup(self, test_mode=False): """Setup the Trainer by calling other setup functions. @@ -119,13 +119,11 @@ class Trainer: """Train the model.""" assert self.pipeline.datamanager.train_input_dataset is not None, "Missing DatsetInputs" - self.visualizer_state.init_scene( - dataset=self.pipeline.datamanager.train_input_dataset, start_train=self.config.viewer.start_train - ) + self._init_viewer_scene() with TimeWriter(writer, EventName.TOTAL_TRAIN_TIME): num_iterations = self.config.trainer.max_num_iterations for step in range(self.start_step, self.start_step + num_iterations): - # if the visualizer used, the rendering of the visualizer will be included in the iteration train time + # if the viewer used, the rendering of the viewer will be included in the iteration train time with TimeWriter(writer, EventName.ITER_TRAIN_TIME, step=step) as train_t: # training callbacks before the training iteration @@ -141,17 +139,9 @@ class Trainer: for callback in self.callbacks: callback.run_callback_at_location(step, location=TrainingCallbackLocation.AFTER_TRAIN_ITERATION) - # NOTE(ethan): we should remove this when the viewer is disabled - with TimeWriter(writer, EventName.ITER_VIS_TIME, step=step) as vis_t: - num_rays_per_batch = self.config.pipeline.datamanager.train_num_rays_per_batch - try: - self.visualizer_state.update_scene(step, self.pipeline.model, num_rays_per_batch) - except RuntimeError: - time.sleep(0.03) # sleep to allow buffer to reset - assert self.visualizer_state.vis is not None - self.visualizer_state.vis["renderingState/log_errors"].write( - "Error: GPU out of memory. Reduce resolution to prevent viewer from crashing." - ) + vis_t = self._update_viewer_state(step) + + self._update_viewer_rays_per_sec(train_t, vis_t, step) # a batch of train rays if step_check(step, self.config.logging.steps_per_log, run_at_zero=True): @@ -183,12 +173,55 @@ class Trainer: if step != 0 and self.config.trainer.steps_per_save and step % self.config.trainer.steps_per_save == 0: self._save_checkpoint(self.config.trainer.model_dir, step) - self._update_rays_per_sec(train_t, vis_t, step) + self._update_viewer_rays_per_sec(train_t, vis_t, step) self._write_out_storage(step) self._write_out_storage(num_iterations) - def _update_rays_per_sec(self, train_t: TimeWriter, vis_t: TimeWriter, step: int): + def _check_viewer_warnings(self) -> None: + """Helper to print out any warnings regarding the way the viewer/loggers are enabled""" + if self.config.viewer.enable: + if self.config.logging.event_writer_config: + string = ( + "[WARNING]: Tensorboard or Wandb enabled with Viewer will slow down Viewer. " + "Please set `--logging.event_writer none` for faster rendering" + ) + CONSOLE.print(f"[bold red]{string}") + string = "[WARNING] Disabling eval iterations since viewer is enabled." + CONSOLE.print(f"[bold red]{string}") + + @check_viewer_enabled + def _init_viewer_scene(self) -> None: + """Initializes viewer scene with given train dataset""" + assert self.viewer_state and self.pipeline.datamanager.train_input_dataset + self.viewer_state.init_scene( + dataset=self.pipeline.datamanager.train_input_dataset, + start_train=self.config.viewer.start_train, + ) + + @check_viewer_enabled + def _update_viewer_state(self, step: int) -> TimeWriter: + """Updates the viewer state by rendering out scene with current pipeline + Returns the time taken to render scene. + + Args: + step: current train step + """ + assert self.viewer_state is not None + with TimeWriter(writer, EventName.ITER_VIS_TIME, step=step) as vis_t: + num_rays_per_batch = self.config.pipeline.datamanager.train_num_rays_per_batch + try: + self.viewer_state.update_scene(step, self.pipeline.model, num_rays_per_batch) + except RuntimeError: + time.sleep(0.03) # sleep to allow buffer to reset + assert self.viewer_state.vis is not None + self.viewer_state.vis["renderingState/log_errors"].write( + "Error: GPU out of memory. Reduce resolution to prevent viewer from crashing." + ) + return vis_t + + @check_viewer_enabled + def _update_viewer_rays_per_sec(self, train_t: TimeWriter, vis_t: TimeWriter, step: int): """Performs update on rays/sec calclation for training Args: diff --git a/nerfactory/pipelines/base.py b/nerfactory/pipelines/base.py index 66addd11b..6b7401a5f 100644 --- a/nerfactory/pipelines/base.py +++ b/nerfactory/pipelines/base.py @@ -63,16 +63,16 @@ class Pipeline(nn.Module): and so on. - TODO: For visualizer functionality to be added down the line, we should make sure - that we are still abstracting away the Model from the end user. The visualizer function + TODO: For viewer functionality to be added down the line, we should make sure + that we are still abstracting away the Model from the end user. The viewer function should probably be done by adding a new iterator on the base data manager that will - talk to the actual visualizer. This is probably ideal to have the visualizer be + talk to the actual viewer. This is probably ideal to have the viewer be primarily located in the data manager (first because it makes sense as the - visualizers main job in this context is to feed in data for the model to load) - so that we can have an easier time ensuring that the visualizer is always + viewers main job in this context is to feed in data for the model to load) + so that we can have an easier time ensuring that the viewer is always returning the same formatted data as for in train / eval. All this is pending changes to be done in the future... but just bear in mind that if learned parameters are in the data manager, - the visualizer may have to use those parameters as well. + the viewer may have to use those parameters as well. Args: diff --git a/nerfactory/utils/decorators.py b/nerfactory/utils/decorators.py index 4558b51bd..77524f699 100644 --- a/nerfactory/utils/decorators.py +++ b/nerfactory/utils/decorators.py @@ -49,12 +49,12 @@ def check_profiler_enabled(func: Callable) -> Callable: return wrapper -def check_visualizer_enabled(func: Callable) -> Callable: - """Decorator: check if visualizer is enabled""" +def check_viewer_enabled(func: Callable) -> Callable: + """Decorator: check if viewer is enabled""" def wrapper(self, *args, **kwargs): ret = None - if self.vis: + if self.config.viewer.enable: ret = func(self, *args, **kwargs) return ret diff --git a/nerfactory/utils/writer.py b/nerfactory/utils/writer.py index 488dac42c..c2ef6c93f 100644 --- a/nerfactory/utils/writer.py +++ b/nerfactory/utils/writer.py @@ -24,8 +24,6 @@ from pathlib import Path from time import time from typing import Any, Dict, List, Optional, Union -import imageio -import numpy as np import torch import wandb from torch.utils.tensorboard import SummaryWriter @@ -47,7 +45,7 @@ class EventName(enum.Enum): ITER_TRAIN_TIME = "Train Iter (time)" TOTAL_TRAIN_TIME = "Train Total (time)" - ITER_VIS_TIME = "Visualizer Rendering (time)" + ITER_VIS_TIME = "Viewer Rendering (time)" ETA = "ETA (time)" TRAIN_RAYS_PER_SEC = "Train Rays / Sec" TEST_RAYS_PER_SEC = "Test Rays / Sec" @@ -157,17 +155,18 @@ def put_time(name: str, duration: float, step: int, avg_over_steps: bool = True, def write_out_storage(): """Function that writes all the events in storage to all the writer locations""" for writer in EVENT_WRITERS: + if isinstance(writer, LocalWriter) and len(EVENT_STORAGE) > 0: + writer.write_stats_log(EVENT_STORAGE[0]["step"]) + continue for event in EVENT_STORAGE: write_func = getattr(writer, event["write_type"].value) write_func(event["name"], event["event"], event["step"]) - # NOTE(ethan): why is the below statement needed? - if isinstance(writer, LocalWriter): - break + EVENT_STORAGE.clear() @check_main_thread -def setup_event_writers(config: cfg.LoggingConfig, max_iter: int, banner_messages: Optional[List[str]] = None) -> None: +def setup_local_writer(config: cfg.LoggingConfig, max_iter: int, banner_messages: Optional[List[str]] = None) -> None: """Initialization of all event writers specified in config Args: @@ -175,14 +174,11 @@ def setup_event_writers(config: cfg.LoggingConfig, max_iter: int, banner_message max_iter: maximum number of train iterations banner_messages: list of messages to always display at bottom of screen """ - for writer_type_config in config.writer: - if writer_type_config.enable: - if isinstance(writer_type_config, cfg.LocalWriterConfig): - curr_writer = writer_type_config.setup(banner_messages=banner_messages) - else: - curr_writer = writer_type_config.setup() - EVENT_WRITERS.append(curr_writer) - logging.info("logging info to: %s", writer_type_config.log_dir) + if config.local_writer.enable: + curr_writer = config.local_writer.setup(banner_messages=banner_messages) + EVENT_WRITERS.append(curr_writer) + else: + logging.info("disabled local writer") ## configure all the global buffer basic information GLOBAL_BUFFER["max_iter"] = max_iter @@ -191,6 +187,22 @@ def setup_event_writers(config: cfg.LoggingConfig, max_iter: int, banner_message GLOBAL_BUFFER["events"] = {} +def setup_event_writer(config: cfg.LoggingConfig) -> None: + """Initialization of all event writers specified in config + + Args: + config: configuration to instantiate loggers + max_iter: maximum number of train iterations + banner_messages: list of messages to always display at bottom of screen + """ + if config.event_writer_config: + curr_writer = config.event_writer_config.setup() + EVENT_WRITERS.append(curr_writer) + logging.info("logging info to: %s", config.event_writer_config.log_dir) + else: + logging.info("disabled tensorboard/wandb event writers") + + class Writer: """Writer class""" @@ -266,7 +278,7 @@ class WandbWriter(Writer): def __init__(self, config: cfg.WandbWriterConfig): super().__init__(config.log_dir) - wandb.init(project="nerfactory-project", dir=config.log_dir, reinit=True) + wandb.init(project="nerfactory-project", dir=str(config.log_dir), reinit=True) def write_image(self, name: str, image: TensorType["H", "W", "C"], step: int) -> None: image = torch.permute(image, (2, 0, 1)) @@ -340,7 +352,7 @@ def _format_time(seconds): @decorate_all([check_main_thread]) -class LocalWriter(Writer): +class LocalWriter: """Local Writer Class TODO: migrate to prettyprint @@ -350,9 +362,8 @@ class LocalWriter(Writer): """ def __init__(self, config: cfg.LocalWriterConfig, banner_messages: Optional[List[str]] = None): - super().__init__(config.log_dir) + self.config = config self.stats_to_track = [name.value for name in config.stats_to_track] - self.max_log_size = config.max_log_size self.keys = set() self.past_mssgs = ["", ""] self.banner_len = 0 if banner_messages is None else len(banner_messages) + 1 @@ -361,19 +372,18 @@ class LocalWriter(Writer): self.past_mssgs.extend(banner_messages) self.has_printed = False - def write_image(self, name: str, image: TensorType["H", "W", "C"], step: int) -> None: - if name in self.stats_to_track and self.log_dir: - image = to8b(image) - image_path = self.log_dir / f"{name}.jpg" - imageio.imwrite(image_path, np.uint8(image.cpu().numpy() * 255.0)) + def write_stats_log(self, step: int) -> None: + """Function to write out scalars to terminal - def write_scalar(self, name: str, scalar: float, step: int) -> None: + Args: + step: current train step + """ if step > 0: - if not self.has_printed and self.max_log_size: + if not self.has_printed and self.config.max_log_size: logging.info( - "\x1b[33;20mPrinting max of %d lines. Set flag `--logging.writer.2.max-log-size=0` " + "\x1b[33;20mPrinting max of %d lines. Set flag `--logging.local_writer.max-log-size=0` " "to disable line wrapping.\x1b[0m", - self.max_log_size, + self.config.max_log_size, ) latest_map, new_key = self._consolidate_events() self._update_header(latest_map, new_key) @@ -405,8 +415,8 @@ class LocalWriter(Writer): latest_map: the most recent dictionary of stats that have been recorded new_key: indicator whether or not there is a new key added to logger """ - full_log_cond = not self.max_log_size and GLOBAL_BUFFER["step"] <= GLOBAL_BUFFER["steps_per_log"] - capped_log_cond = self.max_log_size and (len(self.past_mssgs) - self.banner_len <= 2 or new_key) + full_log_cond = not self.config.max_log_size and GLOBAL_BUFFER["step"] <= GLOBAL_BUFFER["steps_per_log"] + capped_log_cond = self.config.max_log_size and (len(self.past_mssgs) - self.banner_len <= 2 or new_key) if full_log_cond or capped_log_cond: mssg = f"{'Step (% Done)':<20}" for name, _ in latest_map.items(): @@ -441,13 +451,13 @@ class LocalWriter(Writer): curr_mssg += f"{v:<20} " # update the history buffer - if self.max_log_size: + if self.config.max_log_size: if not self.has_printed: cursor_idx = len(self.past_mssgs) - self.banner_len self.has_printed = True else: cursor_idx = len(self.past_mssgs) - if len(self.past_mssgs[2:]) - self.banner_len >= self.max_log_size: + if len(self.past_mssgs[2:]) - self.banner_len >= self.config.max_log_size: self.past_mssgs.pop(2) self.past_mssgs.insert(len(self.past_mssgs) - self.banner_len, curr_mssg) _cursorup(cursor_idx) diff --git a/nerfactory/viewer/server/viewer_utils.py b/nerfactory/viewer/server/viewer_utils.py index 2b3d353f8..6893ad68a 100644 --- a/nerfactory/viewer/server/viewer_utils.py +++ b/nerfactory/viewer/server/viewer_utils.py @@ -12,11 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Code to interface with the `vis/` (the JS visualizer). +"""Code to interface with the `vis/` (the JS viewer). """ +from __future__ import annotations import enum -import logging import os import sys import threading @@ -34,7 +34,6 @@ from nerfactory.configs import base as cfg from nerfactory.datamanagers.datasets import InputDataset from nerfactory.models.base import Model from nerfactory.utils import profiler, visualization, writer -from nerfactory.utils.decorators import check_visualizer_enabled, decorate_all from nerfactory.utils.io import load_from_json, write_to_json from nerfactory.utils.misc import get_dict_to_torch from nerfactory.utils.writer import GLOBAL_BUFFER, EventName, TimeWriter @@ -45,6 +44,19 @@ from nerfactory.viewer.server.visualizer import Viewer console = Console(width=120) +def setup_viewer(config: cfg.ViewerConfig): + """Sets up the viewer if enabled + + Args: + config: the configuration to instantiate viewer + """ + if config.enable: + viewer_state = ViewerState(config) + banner_messages = [f"Viewer at: {viewer_state.viewer_url}"] + return viewer_state, banner_messages + return None, None + + class OutputTypes(str, enum.Enum): """Noncomprehsnive list of output render types""" @@ -67,7 +79,7 @@ class ColormapTypes(str, enum.Enum): class IOChangeException(Exception): - """Basic camera exception to interrupt visualizer""" + """Basic camera exception to interrupt viewer""" class SetTrace: @@ -88,12 +100,12 @@ class RenderThread(threading.Thread): """Thread that does all the rendering calls while listening for interrupts Args: - state: current visualizer state object + state: current viewer state object graph: current checkpoint of model camera_ray_bundle: input rays to pass through the graph to render out """ - def __init__(self, state: "VisualizerState", graph: Model, camera_ray_bundle: RayBundle): + def __init__(self, state: "ViewerState", graph: Model, camera_ray_bundle: RayBundle): threading.Thread.__init__(self) self.state = state self.graph = graph @@ -131,7 +143,7 @@ class CheckThread(threading.Thread): """Thread the constantly checks for io changes and sets a flag indicating interrupt Args: - state: current visualizer state object + state: current viewer state object """ def __init__(self, state): @@ -141,7 +153,7 @@ class CheckThread(threading.Thread): def run(self): """Run function that checks to see if any of the existing state has changed (e.g. camera pose/output type/resolutions). - Sets the visualizer state flag to true to signal + Sets the viewer state flag to true to signal to render thread that an interrupt was registered. """ self.state.check_done_render = False @@ -183,49 +195,41 @@ class CheckThread(threading.Thread): return -@decorate_all([check_visualizer_enabled]) -class VisualizerState: - """Class to hold state for visualizer variables +class ViewerState: + """Class to hold state for viewer variables Args: config: viewer setup configuration """ - def __init__(self, config: cfg.ViewerConfig, config_base_dir: Optional[Path] = None): + def __init__(self, config: cfg.ViewerConfig): self.config = config - self.config_base_dir = config_base_dir - self.vis = None self.viewer_url = None - if self.config.enable: - if self.config.launch_bridge_server: - # start the viewer bridge server - zmq_port = int(self.config.zmq_url.split(":")[-1]) - websocket_port = self.config.websocket_port - self.config.log_filename.parent.mkdir(exist_ok=True) - run_viewer_bridge_server_as_subprocess( - zmq_port, websocket_port, log_filename=str(self.config.log_filename) - ) - # TODO(ethan): move this into the writer such that it's at the bottom - # of the logging stack and easy to see and click - # TODO(ethan): log the output of the viewer bridge server in a file where the training logs go - console.line() - json_filename = os.path.join(os.path.dirname(__file__), "../app/package.json") - version = load_from_json(Path(json_filename))["version"] - self.viewer_url = f"https://viewer.nerfactory.com/{version}/?websocket_url=localhost:{websocket_port}" - viewer_url_local = f"http://localhost:4000/?websocket_url=localhost:{websocket_port}" - pub_open_viewer_instructions_string = f"[Public] Open the viewer at {self.viewer_url}" - dev_open_viewer_instructions_string = f"[Local] Open the viewer at {viewer_url_local}" - console.rule(characters="=") - console.print(pub_open_viewer_instructions_string) - console.print(dev_open_viewer_instructions_string) - console.rule(characters="=") - console.line() - self.vis = Viewer(zmq_url=self.config.zmq_url) - else: - logging.info("Continuing without viewer.") + if self.config.launch_bridge_server: + # start the viewer bridge server + zmq_port = int(self.config.zmq_url.split(":")[-1]) + websocket_port = self.config.websocket_port + self.config.log_filename.parent.mkdir(exist_ok=True) + run_viewer_bridge_server_as_subprocess(zmq_port, websocket_port, log_filename=str(self.config.log_filename)) + # TODO(ethan): move this into the writer such that it's at the bottom + # of the logging stack and easy to see and click + # TODO(ethan): log the output of the viewer bridge server in a file where the training logs go + console.line() + json_filename = os.path.join(os.path.dirname(__file__), "../app/package.json") + version = load_from_json(Path(json_filename))["version"] + self.viewer_url = f"https://viewer.nerfactory.com/{version}/?websocket_url=localhost:{websocket_port}" + viewer_url_local = f"http://localhost:4000/?websocket_url=localhost:{websocket_port}" + pub_open_viewer_instructions_string = f"[Public] Open the viewer at {self.viewer_url}" + dev_open_viewer_instructions_string = f"[Local] Open the viewer at {viewer_url_local}" + console.rule(characters="=") + console.print(pub_open_viewer_instructions_string) + console.print(dev_open_viewer_instructions_string) + console.rule(characters="=") + console.line() + self.vis = Viewer(zmq_url=self.config.zmq_url) - # visualizer specific variables + # viewer specific variables self.prev_camera_matrix = None self.prev_output_type = OutputTypes.INIT self.prev_colormap_type = ColormapTypes.INIT @@ -251,8 +255,7 @@ class VisualizerState: if False, only displays dataset until resume train is toggled """ # set the config base dir - if self.config_base_dir: - self.vis["renderingState/config_base_dir"].write(str(self.config_base_dir)) + self.vis["renderingState/config_base_dir"].write(str(self.config.log_filename.parents[0])) # clear the current scene self.vis["sceneState/sceneBounds"].delete() diff --git a/nerfactory/viewer/server/visualizer.py b/nerfactory/viewer/server/visualizer.py index 41b667de4..dced16d8a 100644 --- a/nerfactory/viewer/server/visualizer.py +++ b/nerfactory/viewer/server/visualizer.py @@ -87,7 +87,7 @@ class ViewerWindow: class Viewer: - """Visualizer class for connecting to the bridge server. + """Viewer class for connecting to the bridge server. Args: zmq_url: Where to connect with ZMQ. Defaults to None. diff --git a/scripts/viewer/view_dataset.py b/scripts/viewer/view_dataset.py index 3ba60ac1e..f3cf6aa28 100755 --- a/scripts/viewer/view_dataset.py +++ b/scripts/viewer/view_dataset.py @@ -4,6 +4,7 @@ view_dataset.py """ import logging +import time from datetime import timedelta import dcargs @@ -15,7 +16,7 @@ from nerfactory.configs import base as cfg from nerfactory.configs.base_configs import AnnotatedBaseConfigUnion from nerfactory.viewer.server import viewer_utils -logging.basicConfig(format="[%(filename)s:%(lineno)d] %(message)s", level=logging.DEBUG) +logging.basicConfig(format="[%(filename)s:%(lineno)d] %(message)s", level=logging.INFO) DEFAULT_TIMEOUT = timedelta(minutes=30) @@ -28,9 +29,11 @@ def main(config: cfg.Config) -> None: if not config.viewer.enable: config.viewer.enable = True logging.info("Enabling viewer to view dataset") - visualizer_state = viewer_utils.VisualizerState(config.viewer) + viewer_state = viewer_utils.ViewerState(config.viewer) datamanager = config.pipeline.datamanager.setup() - visualizer_state.init_scene(dataset=datamanager.train_input_dataset, start_train=False) + viewer_state.init_scene(dataset=datamanager.train_input_dataset, start_train=False) + logging.info("Please refresh and load page at: %s", viewer_state.viewer_url) + time.sleep(30) # allowing time to refresh page if __name__ == "__main__": diff --git a/tests/test_run_train.py b/tests/test_run_train.py index ac203ed43..95c947cf6 100644 --- a/tests/test_run_train.py +++ b/tests/test_run_train.py @@ -31,15 +31,12 @@ def set_reduced_config(config: cfg.Config): ) # use tensorboard logging instead of wandb - config.logging.writer[0].enable = True - config.logging.writer[1].enable = False + config.logging.event_writer = "tb" + config.logging.relative_log_dir = Path("/tmp/") # reduce model factors config.pipeline.model.num_coarse_samples = 4 config.pipeline.model.num_importance_samples = 4 - # set logging to tmp - for writer in config.logging.writer: - writer.log_dir = Path("/tmp/") # remove viewer config.viewer.enable = False @@ -47,6 +44,7 @@ def set_reduced_config(config: cfg.Config): if config.method_name == "instant_ngp": config.pipeline.model.field_implementation = "torch" + config.populate_dynamic_fields() return config