From 2cca012ad87be2568b12ecb005dff0962fb6ab3e Mon Sep 17 00:00:00 2001 From: Ethan Weber Date: Mon, 29 Aug 2022 22:50:57 -0700 Subject: [PATCH] Clickable viewer deployment (#308) * first commit for clickable web viewer * pylint ignore * functional clickable viewer with updated docs * fixed logging --- README.md | 44 ++++++++------- nerfactory/configs/base.py | 2 + .../app/src/modules/WebSocket/WebSocket.jsx | 27 +++++---- nerfactory/viewer/server/server.py | 41 +++++++++++--- nerfactory/viewer/server/subprocess.py | 56 +++++++++++++++++++ nerfactory/viewer/server/viewer_utils.py | 32 ++++++++--- pyproject.toml | 3 + scripts/run_viewer_bridge_server.py | 21 +------ tests/viewer/test_viewer.py | 23 ++++++++ 9 files changed, 181 insertions(+), 68 deletions(-) create mode 100644 nerfactory/viewer/server/subprocess.py create mode 100644 tests/viewer/test_viewer.py diff --git a/README.md b/README.md index af3b2a0e0..18324ec52 100644 --- a/README.md +++ b/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 -``` -
@@ -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//`. -``` +```bash tensorboard --logdir outputs - -# or the following -export TENSORBOARD_PORT= -bash environment/run_tensorboard.sh ```
diff --git a/nerfactory/configs/base.py b/nerfactory/configs/base.py index f40df99b8..6daef9911 100644 --- a/nerfactory/configs/base.py +++ b/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 diff --git a/nerfactory/viewer/app/src/modules/WebSocket/WebSocket.jsx b/nerfactory/viewer/app/src/modules/WebSocket/WebSocket.jsx index 545dc559d..851e85b55 100644 --- a/nerfactory/viewer/app/src/modules/WebSocket/WebSocket.jsx +++ b/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 "?websocket_url=localhost:". 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://" + const websocketEndpoint = getWebsocketEndpoint(); const connect = () => { socket = new WebSocket(`ws://${websocketEndpoint}/`); diff --git a/nerfactory/viewer/server/server.py b/nerfactory/viewer/server/server.py index 312e1939d..1c473006e 100644 --- a/nerfactory/viewer/server/server.py +++ b/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() diff --git a/nerfactory/viewer/server/subprocess.py b/nerfactory/viewer/server/subprocess.py new file mode 100644 index 000000000..0eaac03ef --- /dev/null +++ b/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) diff --git a/nerfactory/viewer/server/viewer_utils.py b/nerfactory/viewer/server/viewer_utils.py index b65a98674..47ee70445 100644 --- a/nerfactory/viewer/server/viewer_utils.py +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 3d1e037ce..15cf377c5 100644 --- a/pyproject.toml +++ b/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 diff --git a/scripts/run_viewer_bridge_server.py b/scripts/run_viewer_bridge_server.py index 709fa9ae0..b31362a79 100644 --- a/scripts/run_viewer_bridge_server.py +++ b/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() diff --git a/tests/viewer/test_viewer.py b/tests/viewer/test_viewer.py new file mode 100644 index 000000000..a556fa44d --- /dev/null +++ b/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.")