Clickable viewer deployment (#308)

* first commit for clickable web viewer

* pylint ignore

* functional clickable viewer with updated docs

* fixed logging

README.md

@@ -64,9 +64,6 @@ pip install -r environment/requirements.txt
 # Install nerfactory as a library
 pip install -e .
 
-# Install library with CUDA support. Change setup.py to `USE_CUDA = True` and then
-python setup.py develop
-
 # Install tiny-cuda-nn (tcnn) and apex to use with the graph_instant_ngp.yaml config
 pip install git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch
 
@@ -95,17 +92,33 @@ Download the original [NeRF dataset](https://drive.google.com/drive/folders/128y
 To run with all the defaults, e.g. vanilla nerf method with the blender lego image
 
 ```
+# To see what models are available.
+python scripts/run_train.py --help
+
 # Run a vanilla nerf model.
-python scripts/run_train.py
+python scripts/run_train.py vanilla_nerf
 
 # Run a faster version with instant ngp using tcnn (without the viewer).
-python scripts/run_train.py --config-name=graph_instant_ngp.yaml
-
-# Run with the viewer. However, you'll have to start the viewer server first. (See the docs.)
-python scripts/run_train.py --config-name=graph_instant_ngp.yaml viewer.enable=true
+python scripts/run_train.py instant_ngp
 ```
 
-With support for [Hydra](https://hydra.cc/), you can run with other configurations by changing appropriate configs defined in `configs/` or by setting flags via command-line arguments.
+#### 3.x Training a model with the viewer
+
+Make sure to forward a port for the websocket to localhost. The default port is 7007, which you should be expose to localhost:7007.
+
+```bash
+# with the default port
+python scripts/run_train.py instant_ngp --viewer.enable
+
+# with a specified websocket port
+python scripts/run_train.py instant_ngp --viewer.enable --viewer.websocket-port=7008
+
+# with the viewer bridge server as a separate process
+# in one terminal, start the bridge server:
+viewer-bridge-server # or equivalently, python scripts/run_viewer_bridge_server.py
+# in another terminal, start training:
+python scripts/run_train.py instant_ngp --viewer.enable --viewer.no-launch-bridge-server
+```
 
 #### 4. Visualizing training runs
 
@@ -116,13 +129,6 @@ We support multiple methods to visualize training, the default configuration use
 
 We have developed our own Real-time web viewer, more information can be found [here](https://plenoptix-nerfactory.readthedocs-hosted.com/en/latest/tooling/viewer.html). This viewer runs during training and is designed to work with models that have fast rendering pipelines.
 
-To enable add the following to your config:
-
-```
-viewer:
-  enable: true
-```
-
 </details>
 
 <details>
@@ -130,12 +136,8 @@ viewer:
 
 If you run everything with the default configuration we log all training curves, test images, and other stats. Once the job is launched, you will be able to track training by launching the tensorboard in `outputs/blender_lego/vanilla_nerf/<timestamp>/<events.tfevents>`.
 
-```
+```bash
 tensorboard --logdir outputs
-
-# or the following
-export TENSORBOARD_PORT=<port>
-bash environment/run_tensorboard.sh
 ```
 
 </details>

nerfactory/configs/base.py

@@ -283,6 +283,8 @@ class ViewerConfig(PrintableConfig):
 
     enable: bool = False
     zmq_url: str = "tcp://127.0.0.1:6000"
+    launch_bridge_server: bool = True
+    websocket_port: int = 7007
     min_render_image_height: int = 64
     max_render_image_height: int = 1024
     num_rays_per_chunk: int = 4096

nerfactory/viewer/app/src/modules/WebSocket/WebSocket.jsx

@@ -10,16 +10,23 @@ const WebSocketContext = createContext(null);
 
 export { WebSocketContext };
 
-function getWebsocketEndpointFromUrl(url: string) {
-  const splitUrl = url.split('?');
-  if (splitUrl.length !== 2) {
-    window.alert("There should be exactly one '?' in the url.");
-    return '';
+function getParam(param_name) {
+  // https://stackoverflow.com/questions/831030/how-to-get-get-request-parameters-in-javascript
+  const params = new RegExp(
+    `[?&]${encodeURIComponent(param_name)}=([^&]*)`,
+  ).exec(window.location.href);
+  if (params === null) {
+    return undefined;
   }
-  const endpoint = splitUrl.pop();
-  if (endpoint === '') {
+  return decodeURIComponent(params[1]);
+}
+
+function getWebsocketEndpoint() {
+  const endpoint = getParam('websocket_url');
+  console.log("websocket url: ", endpoint);
+  if (endpoint === undefined) {
     const message =
-      'Please set the websocket endpoint. E.g., a correct URL may be: http://localhost:4000?localhost:8051';
+      'Please set the websocket endpoint. The format should look like "<viewer_url>?websocket_url=localhost:<port>". E.g., "https://viewer.nerfactory.com/branch/master/?websocket_url=localhost:7007".';
     window.alert(message);
     return '';
   }
@@ -32,8 +39,8 @@ export default function WebSocketContextFunction({ children }) {
   let socket;
   let ws;
 
-  // should look like ws://localhost:8051
-  const websocketEndpoint = getWebsocketEndpointFromUrl(window.location.href);
+  // should look like e.g., "ws://<localhost:port>"
+  const websocketEndpoint = getWebsocketEndpoint();
 
   const connect = () => {
     socket = new WebSocket(`ws://${websocketEndpoint}/`);

nerfactory/viewer/server/server.py

@@ -15,6 +15,7 @@
 """Server bridge to faciliate interactions between python backend and javascript front end"""
 
 
+import argparse
 import sys
 from typing import Callable, List, Optional, Tuple
 
@@ -164,31 +165,39 @@ class ZMQWebSocketBridge:
     """ZMQ web socket bridge class
 
     Args:
-        zmq_url: zmq url to connect to. Defaults to None.
+        zmq_port: zmq port to connect to. Defaults to None.
         host: host of server. Defaults to "127.0.0.1".
         websocket_port: websocket port to connect to. Defaults to None.
     """
 
     context = zmq.Context()
 
-    def __init__(self, zmq_url: Optional[str] = None, host: str = "127.0.0.1", websocket_port: Optional[int] = None):
+    def __init__(
+        self,
+        zmq_port: Optional[int] = None,
+        host: str = "127.0.0.1",
+        websocket_address: str = "0.0.0.0",
+        websocket_port: Optional[int] = None,
+    ):
         self.host = host
+        self.websocket_address = websocket_address
         self.websocket_pool = set()
         self.app = self.make_app()
         self.ioloop = tornado.ioloop.IOLoop.current()
         self.pcs = set()
         self.video_tracks = set()
 
-        if zmq_url is None:
+        if zmq_port is None:
 
             def f(port):
                 return self.setup_zmq(f"{DEFAULT_ZMQ_METHOD:s}://{self.host:s}:{port:d}")
 
             (self.zmq_socket, self.zmq_stream, self.zmq_url), _ = find_available_port(f, DEFAULT_ZMQ_PORT)
         else:
+            zmq_url = f"{DEFAULT_ZMQ_METHOD:s}://{self.host:s}:{zmq_port:d}"
             self.zmq_socket, self.zmq_stream, self.zmq_url = self.setup_zmq(zmq_url)
 
-        listen_kwargs = {}
+        listen_kwargs = {"address": self.websocket_address}
 
         if websocket_port is None:
             _, self.websocket_port = find_available_port(self.app.listen, DEFAULT_WEBSOCKET_PORT, **listen_kwargs)
@@ -196,11 +205,12 @@ class ZMQWebSocketBridge:
             self.app.listen(websocket_port, **listen_kwargs)
             self.websocket_port = websocket_port
 
+        self.websocket_url = f"{self.websocket_address}:{self.websocket_port}"
         self.state_tree = get_tree(StateNode)
 
     def __str__(self) -> str:
         class_name = self.__class__.__name__
-        return f"{class_name} using zmq_url={self.zmq_url} and websocket_port={self.websocket_port}"
+        return f'{class_name} using zmq_url="{self.zmq_url}" and websocket_url="{self.websocket_url}"'
 
     def make_app(self):
         """Create a tornado application for the websocket server."""
@@ -273,7 +283,7 @@ class ZMQWebSocketBridge:
         Args:
             websocket: websocket to send information over
         """
-        print("Sending entire scene state on websocket connect (i.e,. a page refresh).")
+        print("Sending entire scene state due to websocket connection established.")
         for path, node in walk("", self.state_tree):
             if node.data is not None:
                 command = {"type": "write", "path": path, "data": node.data}
@@ -284,6 +294,19 @@ class ZMQWebSocketBridge:
         self.ioloop.start()
 
 
-def start_server_as_subprocess(zmq_url=None):
-    """Starts the ZMQWebSocketBridge server as a subprocess."""
-    raise NotImplementedError()
+def run_viewer_bridge_server():
+    """Run the viewer bridge server"""
+    parser = argparse.ArgumentParser(description="Listen for ZeroMQ commands")
+    parser.add_argument("--zmq-port", "-z", type=int, nargs="?", default=6000)
+    parser.add_argument("--websocket-port", "-wp", type=int, nargs="?", default=7007)
+    args = parser.parse_args()
+    bridge = ZMQWebSocketBridge(zmq_port=args.zmq_port, websocket_port=args.websocket_port)
+    print(bridge)
+    try:
+        bridge.run()
+    except KeyboardInterrupt:
+        pass
+
+
+if __name__ == "__main__":
+    run_viewer_bridge_server()

nerfactory/viewer/server/subprocess.py

@@ -0,0 +1,56 @@
+# Copyright 2022 The Plenoptix Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Code to use the viewer as a subprocess.
+"""
+
+import atexit
+import os
+import subprocess
+import sys
+from typing import Optional
+
+from nerfactory.viewer.server import server
+
+
+def run_viewer_bridge_server_as_subprocess(zmq_port: int, websocket_port: int, log_filename: Optional[str] = None):
+    """Runs the viewer bridge server as a subprocess.
+
+    Args:
+        zmq_port: Port to use for the ZMQ server.
+        websocket_port: Port to use for the websocket server.
+        log_filename: Filename to use for the log file. Defaults to None. If None, no log file is created.
+
+    Returns:
+        None
+    """
+    args = [sys.executable, "-u", "-m", server.__name__]
+    args.append("--zmq-port")
+    args.append(str(zmq_port))
+    args.append("--websocket-port")
+    args.append(str(websocket_port))
+    # supress output if no log filename is specified
+    logfile = open(  # pylint: disable=consider-using-with
+        log_filename if log_filename else os.devnull, "w", encoding="utf8"
+    )
+    process = subprocess.Popen(  # pylint: disable=consider-using-with
+        args, stdout=logfile, stderr=logfile, start_new_session=True
+    )
+
+    def cleanup(process):
+        process.kill()
+        process.wait()
+
+    atexit.register(cleanup, process)

nerfactory/viewer/server/viewer_utils.py

@@ -34,6 +34,7 @@ from nerfactory.models.base import Model
 from nerfactory.utils import profiler, visualization
 from nerfactory.utils.decorators import check_visualizer_enabled, decorate_all
 from nerfactory.utils.writer import GLOBAL_BUFFER, EventName
+from nerfactory.viewer.server.subprocess import run_viewer_bridge_server_as_subprocess
 from nerfactory.viewer.server.utils import get_intrinsics_matrix_and_camera_to_world_h
 from nerfactory.viewer.server.visualizer import Viewer
 
@@ -188,9 +189,29 @@ class VisualizerState:
         self.config = config
 
         self.vis = None
+        self.viewer_url = None
         if self.config.enable:
-            zmq_url = self.config.zmq_url
-            self.vis = Viewer(zmq_url=zmq_url)
+            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
+                run_viewer_bridge_server_as_subprocess(zmq_port, websocket_port)
+                # 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
+                print("\n")
+                self.viewer_url = (
+                    f"https://viewer.nerfactory.com/branch/master/?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}"'
+                print("-" * len(pub_open_viewer_instructions_string))
+                print(pub_open_viewer_instructions_string)
+                print(dev_open_viewer_instructions_string)
+                print("-" * len(pub_open_viewer_instructions_string))
+                print("\n")
+            self.vis = Viewer(zmq_url=self.config.zmq_url)
         else:
             logging.info("Continuing without viewer.")
 
@@ -485,10 +506,3 @@ class VisualizerState:
             stuff_colors = graph.stuff_colors if hasattr(graph, "stuff_colors") else None
             self._send_output_to_viewer(outputs, stuff_colors=stuff_colors)
             self._update_viewer_stats(render_duration, image_height)
-
-
-def get_default_vis() -> Viewer:
-    """Returns the default Visualizer."""
-    zmq_url = "tcp://0.0.0.0:6000"
-    viewer = Viewer(zmq_url=zmq_url)
-    return viewer

pyproject.toml

@@ -17,6 +17,9 @@ classifiers = [
 [project.urls]
 "Documentation" = "https://plenoptix-nerfactory.readthedocs-hosted.com/en/latest/?badge=latest"
 
+[project.scripts]
+viewer-bridge-server = "nerfactory.viewer.server.server:run_viewer_bridge_server"
+
 # black
 [tool.black]
 line-length = 120

scripts/run_viewer_bridge_server.py

@@ -1,23 +1,6 @@
 """View Bridge Server"""
 
-import argparse
-
-from nerfactory.viewer.server.server import ZMQWebSocketBridge
-
-
-def main():
-    """Run the viewer bridge server"""
-    parser = argparse.ArgumentParser(description="Listen for ZeroMQ commands")
-    parser.add_argument("--zmq-url", "-z", type=str, nargs="?", default=None)
-    parser.add_argument("--websocket-port", "-wp", type=str, nargs="?", default=None)
-    args = parser.parse_args()
-    bridge = ZMQWebSocketBridge(zmq_url=args.zmq_url, websocket_port=args.websocket_port)
-    print(bridge)
-    try:
-        bridge.run()
-    except KeyboardInterrupt:
-        pass
-
+from nerfactory.viewer.server.server import run_viewer_bridge_server
 
 if __name__ == "__main__":
-    main()
+    run_viewer_bridge_server()

tests/viewer/test_viewer.py

@@ -0,0 +1,23 @@
+"""
+Test communications with the viewer.
+"""
+
+import time
+
+from nerfactory.viewer.server.subprocess import run_viewer_bridge_server_as_subprocess
+
+
+def test_run_subprocess():
+    """Test running the viewer bridge server as a subprocess."""
+    zmq_port = 6000 + 10  # don't use the default ports to avoid conflicting with current training jobs
+    websocket_port = 7007 + 10
+    print("Starting the viewer bridge server as a subprocess.")
+    run_viewer_bridge_server_as_subprocess(zmq_port, websocket_port)
+    print("Started the viewer bridge server as a subprocess.")
+    # do something will the viewer bridge server runs
+    time.sleep(1)
+
+
+if __name__ == "__main__":
+    test_run_subprocess()
+    print("Tests passed.")