Add trainer parsing (#1801)

* Increase max_res of first proposal network

* add trainer to parsing for ViewerElements

* lint

---------

Co-authored-by: Ayaan Haque <ayaanzhaque@gmail.com>

nerfstudio/viewer/server/gui_utils.py

@@ -20,9 +20,6 @@ from typing import Any, List, Tuple
 
 from torch import nn
 
-from nerfstudio.pipelines.base_pipeline import Pipeline
-from nerfstudio.viewer.server.viewer_elements import ViewerElement
-
 
 def parse_object(
     obj: Any,
@@ -68,14 +65,3 @@ def parse_object(
             for ts, o in lower_rets:
                 add(ret, ts, o)
     return ret
-
-
-def get_viewer_elements(pipeline: Pipeline) -> List[Tuple[str, ViewerElement]]:
-    """
-    Recursively parse the pipeline object and return a tree of all the ViewerElements contained
-
-    returns a list of (path/to/object, param), which represents the path down the object tree as well as
-    the ViewerElement itself
-    """
-    ret = parse_object(pipeline, ViewerElement, "Custom")
-    return ret

nerfstudio/viewer/server/viewer_state.py

@@ -37,7 +37,7 @@ from nerfstudio.utils.io import load_from_json, write_to_json
 from nerfstudio.utils.writer import GLOBAL_BUFFER, EventName
 from nerfstudio.viewer.server import viewer_utils
 from nerfstudio.viewer.server.control_panel import ControlPanel
-from nerfstudio.viewer.server.gui_utils import get_viewer_elements
+from nerfstudio.viewer.server.gui_utils import parse_object
 from nerfstudio.viewer.server.render_state_machine import (
     RenderAction,
     RenderStateMachine,
@@ -149,7 +149,11 @@ class ViewerState:
                 with self.viser_server.gui_folder(folder_labels[0]):
                     nested_folder_install(folder_labels[1:], element)
 
-        self.viewer_elements = get_viewer_elements(self.pipeline)
+        self.viewer_elements = []
+        if self.trainer is not None:
+            self.viewer_elements.extend(parse_object(self.trainer, ViewerElement, "Trainer"))
+
+        self.viewer_elements.extend(parse_object(pipeline, ViewerElement, "Pipeline"))
         for param_path, element in self.viewer_elements:
             folder_labels = param_path.split("/")[:-1]
             nested_folder_install(folder_labels, element)