From e8a5941b241a5993778f90ddb939d162a5549604 Mon Sep 17 00:00:00 2001
From: Brent Yi <yibrenth@gmail.com>
Date: Wed, 28 Sep 2022 12:25:10 -0700
Subject: [PATCH] Don't require model type for `view_dataset.py` (#568)

---
 nerfactory/datamanagers/datasets.py |  2 +-
 scripts/viewer/view_dataset.py      | 27 +++++++++------------------
 2 files changed, 10 insertions(+), 19 deletions(-)
 mode change 100755 => 100644 scripts/viewer/view_dataset.py

diff --git a/nerfactory/datamanagers/datasets.py b/nerfactory/datamanagers/datasets.py
index 2404f0fac..5f7214302 100644
--- a/nerfactory/datamanagers/datasets.py
+++ b/nerfactory/datamanagers/datasets.py
@@ -34,7 +34,7 @@ class InputDataset(Dataset):
     """Dataset that returns images.
 
     Args:
-        config: the DataParserConfig used to instantiate class, which in turn is used to load dataset inputs
+        dataset_inputs: description of where and how to read input images.
     """
 
     def __init__(self, dataset_inputs: DatasetInputs):
diff --git a/scripts/viewer/view_dataset.py b/scripts/viewer/view_dataset.py
old mode 100755
new mode 100644
index 682db3b22..677aff2d9
--- a/scripts/viewer/view_dataset.py
+++ b/scripts/viewer/view_dataset.py
@@ -6,14 +6,14 @@ view_dataset.py
 import logging
 import time
 from datetime import timedelta
+from pathlib import Path
 
 import dcargs
 import torch
-import yaml
-from rich.console import Console
 
 from nerfactory.configs import base as cfg
-from nerfactory.configs.base_configs import AnnotatedBaseConfigUnion
+from nerfactory.datamanagers.base import AnnotatedDataParserUnion
+from nerfactory.datamanagers.datasets import InputDataset
 from nerfactory.viewer.server import viewer_utils
 
 logging.basicConfig(format="[%(filename)s:%(lineno)d] %(message)s", level=logging.INFO)
@@ -24,11 +24,12 @@ DEFAULT_TIMEOUT = timedelta(minutes=30)
 torch.backends.cudnn.benchmark = True  # type: ignore
 
 
-def main(config: cfg.Config) -> None:
+def main(
+    dataparser: AnnotatedDataParserUnion,
+    viewer: cfg.ViewerConfig,
+    log_base_dir: Path = Path("/tmp/nerfactory_viewer_logs"),
+) -> None:
     """Main function."""
-    if not config.is_viewer_enabled():
-        config.vis.append("viewer")
-        logging.info("Enabling viewer to view dataset")
     viewer_state = viewer_utils.ViewerState(config.viewer)
     datamanager = config.pipeline.datamanager.setup()
     viewer_state.init_scene(dataset=datamanager.train_input_dataset, start_train=False)
@@ -37,15 +38,5 @@ def main(config: cfg.Config) -> None:
 
 
 if __name__ == "__main__":
-    console = Console(width=120)
-
     dcargs.extras.set_accent_color("bright_yellow")
-    instantiated_config = dcargs.cli(AnnotatedBaseConfigUnion)
-    if instantiated_config.trainer.load_config:
-        logging.info(f"Loading pre-set config from: {instantiated_config.trainer.load_config}")
-        instantiated_config = yaml.load(instantiated_config.trainer.load_config.read_text(), Loader=yaml.Loader)
-
-    console.rule("Config")
-    console.print(instantiated_config)
-    console.rule("")
-    main(config=instantiated_config)
+    dcargs.cli(main)