From d71c05f2560061e36a33c9de7115b1a080c9b6a3 Mon Sep 17 00:00:00 2001 From: Ethan Weber Date: Fri, 14 Oct 2022 15:05:28 -0700 Subject: [PATCH] Google Colab support with viewer (#747) * allows bridge server to run on a different machine * added script * update version * starting docs * update docs and scripts * ngrok support * use tcp * moving webrtc outside bridge server. wop * lazy commit * lazy commit * running code. webrtc not working * lazy commit * better logic for webrtc connection * cleanup * todo comment * working webrtc * removed logging * fixing black * fix eslint * fix python version * fix readme * misc cleanup * small docs fix Co-authored-by: Ethan Weber --- README.md | 13 ++- docs/index.md | 5 +- docs/quickstart/google_colab.md | 3 + nerfstudio/configs/base_config.py | 17 ++-- nerfstudio/configs/config_utils.py | 6 +- .../dataparsers/instant_ngp_dataparser.py | 6 +- .../data/dataparsers/nerfstudio_dataparser.py | 8 +- nerfstudio/engine/trainer.py | 23 ++--- nerfstudio/utils/eval_utils.py | 12 +-- nerfstudio/utils/install_checks.py | 6 +- nerfstudio/utils/printing.py | 12 +-- nerfstudio/utils/profiler.py | 9 +- nerfstudio/utils/writer.py | 17 ++-- nerfstudio/viewer/app/package.json | 2 +- .../viewer/app/src/modules/Banner/Banner.jsx | 15 ++- .../src/modules/LandingModal/LandingModal.jsx | 18 ++-- .../SidePanel/StatusPanel/StatusPanel.jsx | 11 +++ .../src/modules/WebRtcWindow/WebRtcWindow.jsx | 10 +- .../app/src/modules/WebSocket/WebSocket.jsx | 14 +-- .../app/src/modules/WebSocketUrlField.jsx | 22 +++-- nerfstudio/viewer/app/src/reducer.js | 2 +- nerfstudio/viewer/server/server.py | 72 +++----------- nerfstudio/viewer/server/utils.py | 16 ++++ nerfstudio/viewer/server/video_stream.py | 55 ----------- nerfstudio/viewer/server/viewer_utils.py | 96 ++++++++++++++++--- nerfstudio/viewer/server/visualizer.py | 46 ++++----- pyproject.toml | 2 + scripts/completions/install.py | 3 + scripts/docs/add_nb_tags.py | 8 +- scripts/docs/build_docs.py | 16 ++-- scripts/downloads/download_data.py | 14 +-- scripts/eval.py | 7 +- scripts/github/run_actions.py | 34 +++---- scripts/render.py | 19 ++-- scripts/train.py | 17 ++-- scripts/viewer/run_bridge_server.py | 7 -- scripts/viewer/view_dataset.py | 10 +- 37 files changed, 324 insertions(+), 329 deletions(-) create mode 100644 docs/quickstart/google_colab.md delete mode 100644 scripts/viewer/run_bridge_server.py diff --git a/README.md b/README.md index ed5bb264f..11239e7ed 100644 --- a/README.md +++ b/README.md @@ -46,7 +46,6 @@ - - [Quickstart](#quickstart) - [Learn more](#learn-more) - [Supported Features](#supported-features) @@ -190,7 +189,7 @@ If you're interested in learning more on how to create your own pipelines, devel | [Documentation](https://docs.nerf.studio/en/latest/) | Full API documentation and tutorials | | [Viewer](https://viewer.nerf.studio/) | Home page for our web viewer | | 🎒 **Educational** | -| [Model Descriptions](https://docs.nerf.studio/en/latest/nerfology/methods/index.html) | Description of all the models supported by nerfstudio and explanations of component parts. | +| [Model Descriptions](https://docs.nerf.studio/en/latest/nerfology/methods/index.html) | Description of all the models supported by nerfstudio and explanations of component parts. | | [Component Descriptions](https://docs.nerf.studio/en/latest/nerfology/model_components/index.html) | Interactive notebooks that explain notable/commonly used modules in various models. | | 🏃 **Tutorials** | | [Getting Started](https://docs.nerf.studio/en/latest/quickstart/installation.html) | A more in-depth guide on how to get started with nerfstudio from installation to contributing. | @@ -200,9 +199,9 @@ If you're interested in learning more on how to create your own pipelines, devel | [Creating datasets](https://docs.nerf.studio/en/latest/quickstart/custom_dataset.html) | Have a new dataset? Learn how to run it with nerfstudio. | | | [Contributing](https://docs.nerf.studio/en/latest/reference/contributing.html) | Walk-through for how you can start contributing now. | | 💖 **Community** | -| [Discord](https://discord.gg/uMbNqcraFc) | Join our community to discuss more. We would love to hear from you! | -| [Twitter](https://twitter.com/nerfstudioteam) | Follow us on Twitter @nerfstudioteam to see cool updates and announcements | -| [TikTok](#) | Coming soon! Follow us on TikTok to see some of our fan favorite results | +| [Discord](https://discord.gg/uMbNqcraFc) | Join our community to discuss more. We would love to hear from you! | +| [Twitter](https://twitter.com/nerfstudioteam) | Follow us on Twitter @nerfstudioteam to see cool updates and announcements | +| [TikTok](#) | Coming soon! Follow us on TikTok to see some of our fan favorite results | # Supported Features @@ -238,8 +237,8 @@ If you use this library or find the documentation useful for your research, plea ``` @misc{nerfstudio, title={Nerfstudio: A Framework for Neural Radiance Field Development}, - author={Matthew Tancik*, Ethan Weber*, Evonne Ng*, Ruilong Li, Brent Yi, - Terrance Wang, Alexander Kristoffersen, Jake Austin, Kamyar Salahi, + author={Matthew Tancik*, Ethan Weber*, Evonne Ng*, Ruilong Li, Brent Yi, + Terrance Wang, Alexander Kristoffersen, Jake Austin, Kamyar Salahi, Abhik Ahuja, David McAllister, Angjoo Kanazawa}, year={2022}, url={https://github.com/nerfstudio-project/nerfstudio}, diff --git a/docs/index.md b/docs/index.md index 1682365ff..c6bbb80c8 100644 --- a/docs/index.md +++ b/docs/index.md @@ -50,6 +50,7 @@ quickstart/installation quickstart/first_nerf quickstart/custom_dataset quickstart/viewer_quickstart +quickstart/google_colab Contributing ``` @@ -126,7 +127,9 @@ If you use this library or find the documentation useful for your research, plea ```none @misc{nerfstudio, title={Nerfstudio: A Framework for Neural Radiance Field Development}, - author={Matthew Tancik*, Ethan Weber*, Evonne Ng*, Ruilong Li, Brent Yi, Terrance Wang, Alexander Kristoffersen, Jake Austin, Kamyar Salahi, Abhik Ahuja, David McAllister, Angjoo Kanazawa}, + author={Matthew Tancik*, Ethan Weber*, Evonne Ng*, Ruilong Li, Brent Yi, + Terrance Wang, Alexander Kristoffersen, Jake Austin, Kamyar Salahi, + Abhik Ahuja, David McAllister, Angjoo Kanazawa}, year={2022}, url={https://github.com/nerfstudio-project/nerfstudio}, } diff --git a/docs/quickstart/google_colab.md b/docs/quickstart/google_colab.md new file mode 100644 index 000000000..b086e5e59 --- /dev/null +++ b/docs/quickstart/google_colab.md @@ -0,0 +1,3 @@ +# Google Colab + +Nerfstudio works in Google Colab using their hosted GPUs! Here is [the link](https://colab.research.google.com/drive/1l-QDYEcv759kUNJQrJ8z1CI8o1iIiU-l?usp=sharing). Make a copy of this file and go from there! diff --git a/nerfstudio/configs/base_config.py b/nerfstudio/configs/base_config.py index 218050a79..cecbc6ffe 100644 --- a/nerfstudio/configs/base_config.py +++ b/nerfstudio/configs/base_config.py @@ -18,7 +18,6 @@ from __future__ import annotations -import logging from dataclasses import dataclass from datetime import datetime from pathlib import Path @@ -32,6 +31,7 @@ from nerfstudio.configs.config_utils import to_immutable_dict # model instances from nerfstudio.utils import writer +CONSOLE = Console(width=120) # Pretty printing class class PrintableConfig: # pylint: disable=too-few-public-methods @@ -167,9 +167,11 @@ class ViewerConfig(PrintableConfig): zmq_port: Optional[int] = None """The zmq port to connect to for communication. If None, find an available port.""" launch_bridge_server: bool = True - """whether or not to launch the zmq bridge server""" - websocket_port: int = 7007 + """whether or not to launch the bridge server""" + websocket_port: Optional[int] = 7007 """the default websocket port to connect to""" + ip_address: str = "0.0.0.0" + """the ip address where the bridge server is running""" num_rays_per_chunk: int = 32768 """number of rays per chunk to render with viewer""" @@ -249,10 +251,9 @@ class Config(PrintableConfig): def print_to_terminal(self) -> None: """Helper to pretty print config to terminal""" - console = Console(width=120) - console.rule("Config") - console.print(self) - console.rule("") + CONSOLE.rule("Config") + CONSOLE.print(self) + CONSOLE.rule("") def save_config(self) -> None: """Save config to base directory""" @@ -260,5 +261,5 @@ class Config(PrintableConfig): assert base_dir is not None base_dir.mkdir(parents=True, exist_ok=True) config_yaml_path = base_dir / "config.yml" - logging.info(f"Saving config to: {config_yaml_path}") + CONSOLE.log(f"Saving config to: {config_yaml_path}") config_yaml_path.write_text(yaml.dump(self), "utf8") diff --git a/nerfstudio/configs/config_utils.py b/nerfstudio/configs/config_utils.py index 259603cb3..59b1b252b 100644 --- a/nerfstudio/configs/config_utils.py +++ b/nerfstudio/configs/config_utils.py @@ -23,6 +23,7 @@ from typing import Any, Dict from rich.console import Console +CONSOLE = Console() # pylint: disable=import-outside-toplevel # cannot use mutable types directly within dataclass; abstracting default factory calls @@ -44,7 +45,6 @@ def convert_markup_to_ansi(markup_string: str) -> str: Returns: Text formatted via ANSI sequences. """ - console = Console() - with console.capture() as out: - console.print(markup_string, soft_wrap=True) + with CONSOLE.capture() as out: + CONSOLE.print(markup_string, soft_wrap=True) return out.get() diff --git a/nerfstudio/data/dataparsers/instant_ngp_dataparser.py b/nerfstudio/data/dataparsers/instant_ngp_dataparser.py index 613b6b6c1..1641ad43d 100644 --- a/nerfstudio/data/dataparsers/instant_ngp_dataparser.py +++ b/nerfstudio/data/dataparsers/instant_ngp_dataparser.py @@ -16,13 +16,13 @@ from __future__ import annotations -import logging from dataclasses import dataclass, field from pathlib import Path from typing import Dict, Tuple, Type import numpy as np import torch +from rich.console import Console from nerfstudio.cameras import camera_utils from nerfstudio.cameras.cameras import Cameras, CameraType @@ -34,6 +34,8 @@ from nerfstudio.data.dataparsers.base_dataparser import ( from nerfstudio.data.scene_box import SceneBox from nerfstudio.utils.io import load_from_json +CONSOLE = Console(width=120) + @dataclass class InstantNGPDataParserConfig(DataParserConfig): @@ -69,7 +71,7 @@ class InstantNGP(DataParser): image_filenames.append(fname) poses.append(np.array(frame["transform_matrix"])) if num_skipped_image_filenames >= 0: - logging.info("Skipping %s files in dataset split %s.", num_skipped_image_filenames, split) + CONSOLE.print(f"Skipping {num_skipped_image_filenames} files in dataset split {split}.") assert ( len(image_filenames) != 0 ), """ diff --git a/nerfstudio/data/dataparsers/nerfstudio_dataparser.py b/nerfstudio/data/dataparsers/nerfstudio_dataparser.py index 35f2412cf..1e29f6b11 100644 --- a/nerfstudio/data/dataparsers/nerfstudio_dataparser.py +++ b/nerfstudio/data/dataparsers/nerfstudio_dataparser.py @@ -15,7 +15,6 @@ from __future__ import annotations -import logging import math from dataclasses import dataclass, field from pathlib import Path, PureWindowsPath @@ -36,8 +35,7 @@ from nerfstudio.data.dataparsers.base_dataparser import ( from nerfstudio.data.scene_box import SceneBox from nerfstudio.utils.io import load_from_json -console = Console() - +CONSOLE = Console(width=120) MAX_AUTO_RESOLUTION = 1600 @@ -88,7 +86,7 @@ class Nerfstudio(DataParser): image_filenames.append(fname) poses.append(np.array(frame["transform_matrix"])) if num_skipped_image_filenames >= 0: - logging.info("Skipping %s files in dataset split %s.", num_skipped_image_filenames, split) + CONSOLE.log(f"Skipping {num_skipped_image_filenames} files in dataset split {split}.") assert ( len(image_filenames) != 0 ), """ @@ -185,8 +183,8 @@ class Nerfstudio(DataParser): break df += 1 - console.print(f"Auto image downscale factor of {2**df}") self.downscale_factor = 2**df + CONSOLE.log(f"Auto image downscale factor of {self.downscale_factor}") else: self.downscale_factor = self.config.downscale_factor diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index 78f7ba0a9..686cfcca4 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -19,13 +19,12 @@ from __future__ import annotations import dataclasses import functools -import logging import os import time from typing import Any, Dict, List, Tuple import torch -from rich import console +from rich.console import Console from torch.cuda.amp.grad_scaler import GradScaler from nerfstudio.configs import base_config as cfg @@ -46,8 +45,7 @@ from nerfstudio.utils.misc import step_check from nerfstudio.utils.writer import EventName, TimeWriter from nerfstudio.viewer.server import viewer_utils -logging.getLogger("PIL").setLevel(logging.WARNING) -CONSOLE = console.Console() +CONSOLE = Console(width=120) def train_loop(local_rank: int, world_size: int, config: cfg.Config) -> Any: @@ -94,7 +92,7 @@ class Trainer: self.mixed_precision = self.config.trainer.mixed_precision if self.device == "cpu": self.mixed_precision = False - logging.warning("Mixed precision is disabled for CPU training.") + CONSOLE.print("Mixed precision is disabled for CPU training.") self._start_step = 0 # optimizers self.grad_scaler = GradScaler(enabled=self.mixed_precision) @@ -102,13 +100,12 @@ class Trainer: self.base_dir = config.get_base_dir() # directory to save checkpoints self.checkpoint_dir = config.get_checkpoint_dir() - logging.info("Saving checkpoints to: %s", self.checkpoint_dir) + CONSOLE.log(f"Saving checkpoints to: {self.checkpoint_dir}") # set up viewer if enabled viewer_log_path = self.base_dir / config.viewer.relative_log_filename - ret = None + self.viewer_state, banner_messages = None, None if self.config.is_viewer_enabled(): - ret = viewer_utils.setup_viewer(config.viewer, log_filename=viewer_log_path) - (self.viewer_state, banner_messages) = ret if ret else (None, None) + self.viewer_state, banner_messages = viewer_utils.setup_viewer(config.viewer, log_filename=viewer_log_path) self._check_viewer_warnings() # set up writers/profilers if enabled writer_log_path = self.base_dir / config.logging.relative_log_dir @@ -145,7 +142,7 @@ class Trainer: """Train the model.""" assert self.pipeline.datamanager.train_dataset is not None, "Missing DatsetInputs" - self._init_viewer_scene() + self._init_viewer_state() 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): @@ -207,7 +204,7 @@ class Trainer: CONSOLE.print(f"[bold red]{string}") @check_viewer_enabled - def _init_viewer_scene(self) -> None: + def _init_viewer_state(self) -> None: """Initializes viewer scene with given train dataset""" assert self.viewer_state and self.pipeline.datamanager.train_dataset self.viewer_state.init_scene( @@ -271,9 +268,9 @@ class Trainer: self.pipeline.load_pipeline(loaded_state["pipeline"]) self.optimizers.load_optimizers(loaded_state["optimizers"]) self.grad_scaler.load_state_dict(loaded_state["scalers"]) - logging.info("done loading checkpoint from %s", load_path) + CONSOLE.print(f"done loading checkpoint from {load_path}") else: - logging.info("No checkpoints to load, training from scratch") + CONSOLE.print("No checkpoints to load, training from scratch") @check_main_thread def save_checkpoint(self, step: int) -> None: diff --git a/nerfstudio/utils/eval_utils.py b/nerfstudio/utils/eval_utils.py index a97997ffa..2f1ee4673 100644 --- a/nerfstudio/utils/eval_utils.py +++ b/nerfstudio/utils/eval_utils.py @@ -29,7 +29,7 @@ from rich.console import Console from nerfstudio.configs import base_config as cfg from nerfstudio.pipelines.base_pipeline import Pipeline -console = Console(width=120) +CONSOLE = Console(width=120) def eval_load_checkpoint(config: cfg.TrainerConfig, pipeline: Pipeline) -> Path: @@ -42,12 +42,12 @@ def eval_load_checkpoint(config: cfg.TrainerConfig, pipeline: Pipeline) -> Path: """ assert config.load_dir is not None if config.load_step is None: - console.print("Loading latest checkpoint from load_dir") + CONSOLE.print("Loading latest checkpoint from load_dir") # NOTE: this is specific to the checkpoint name format if not os.path.exists(config.load_dir): - console.rule("Error", style="red") - console.print(f"No checkpoint directory found at {config.load_dir}, ", justify="center") - console.print( + CONSOLE.rule("Error", style="red") + CONSOLE.print(f"No checkpoint directory found at {config.load_dir}, ", justify="center") + CONSOLE.print( "Please make sure the checkpoint exists, they should be generated periodically during training", justify="center", ) @@ -59,7 +59,7 @@ def eval_load_checkpoint(config: cfg.TrainerConfig, pipeline: Pipeline) -> Path: assert load_path.exists(), f"Checkpoint {load_path} does not exist" loaded_state = torch.load(load_path, map_location="cpu") pipeline.load_pipeline(loaded_state["pipeline"]) - console.print(f":white_check_mark: Done loading checkpoint from {load_path}") + CONSOLE.print(f":white_check_mark: Done loading checkpoint from {load_path}") return load_path diff --git a/nerfstudio/utils/install_checks.py b/nerfstudio/utils/install_checks.py index 0372130a8..23351a1b2 100644 --- a/nerfstudio/utils/install_checks.py +++ b/nerfstudio/utils/install_checks.py @@ -19,14 +19,14 @@ import sys from rich.console import Console -console = Console(width=120) +CONSOLE = Console(width=120) def check_ffmpeg_installed(): """Checks if ffmpeg is installed.""" ffmpeg_path = shutil.which("ffmpeg") if ffmpeg_path is None: - console.print("[bold red]Could not find ffmpeg. Please install ffmpeg.") + CONSOLE.print("[bold red]Could not find ffmpeg. Please install ffmpeg.") print("See https://ffmpeg.org/download.html for installation instructions.") print("ffmpeg is only necessary if using videos as input.") sys.exit(1) @@ -36,6 +36,6 @@ def check_colmap_installed(): """Checks if colmap is installed.""" colmap_path = shutil.which("colmap") if colmap_path is None: - console.print("[bold red]Could not find COLMAP. Please install COLMAP.") + CONSOLE.print("[bold red]Could not find COLMAP. Please install COLMAP.") print("See https://colmap.github.io/install.html for installation instructions.") sys.exit(1) diff --git a/nerfstudio/utils/printing.py b/nerfstudio/utils/printing.py index ba2556cf7..e6ebc18d1 100644 --- a/nerfstudio/utils/printing.py +++ b/nerfstudio/utils/printing.py @@ -18,19 +18,19 @@ from math import floor, log from rich.console import Console -console = Console(width=120) +CONSOLE = Console(width=120) def print_tcnn_speed_warning(method_name: str): """Prints a warning about the speed of the TCNN.""" - console.line() - console.print(f"[bold yellow]WARNING: Using a slow implementation of {method_name}. ") - console.print( + CONSOLE.line() + CONSOLE.print(f"[bold yellow]WARNING: Using a slow implementation of {method_name}. ") + CONSOLE.print( "[bold yellow]:person_running: :person_running: " + "Install tcnn for speedups :person_running: :person_running:" ) - console.print("[yellow]pip install git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch") - console.line() + CONSOLE.print("[yellow]pip install git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch") + CONSOLE.line() def human_format(num): diff --git a/nerfstudio/utils/profiler.py b/nerfstudio/utils/profiler.py index 3c90ef791..c1fc96dcb 100644 --- a/nerfstudio/utils/profiler.py +++ b/nerfstudio/utils/profiler.py @@ -17,10 +17,11 @@ Profiler base class and functionality """ from __future__ import annotations -import logging import time from typing import Callable +from rich.console import Console + from nerfstudio.configs import base_config as cfg from nerfstudio.utils import comms from nerfstudio.utils.decorators import ( @@ -29,6 +30,8 @@ from nerfstudio.utils.decorators import ( decorate_all, ) +CONSOLE = Console(width=120) + PROFILER = [] @@ -82,7 +85,7 @@ class Profiler: def print_profile(self): """helper to print out the profiler stats""" - logging.info("Printing profiling stats, from longest to shortest duration in seconds") + CONSOLE.print("Printing profiling stats, from longest to shortest duration in seconds") sorted_keys = sorted( self.profiler_dict.keys(), key=lambda k: self.profiler_dict[k]["val"], @@ -90,4 +93,4 @@ class Profiler: ) for k in sorted_keys: val = f"{self.profiler_dict[k]['val']:0.4f}" - print(f"{k:<20}: {val:<20}") + CONSOLE.print(f"{k:<20}: {val:<20}") diff --git a/nerfstudio/utils/writer.py b/nerfstudio/utils/writer.py index d7022b7fa..bd9741a7d 100644 --- a/nerfstudio/utils/writer.py +++ b/nerfstudio/utils/writer.py @@ -18,7 +18,6 @@ Generic Writer class from __future__ import annotations import enum -import logging from abc import abstractmethod from pathlib import Path from time import time @@ -26,7 +25,7 @@ from typing import Any, Dict, List, Optional, Union import torch import wandb -from rich import console +from rich.console import Console from torch.utils.tensorboard import SummaryWriter from torchtyping import TensorType @@ -34,8 +33,7 @@ from nerfstudio.configs import base_config as cfg from nerfstudio.utils.decorators import check_main_thread, decorate_all from nerfstudio.utils.printing import human_format -CONSOLE = console.Console() - +CONSOLE = Console(width=120) to8b = lambda x: (255 * torch.clamp(x, min=0, max=1)).to(torch.uint8) EVENT_WRITERS = [] EVENT_STORAGE = [] @@ -181,7 +179,7 @@ def setup_local_writer(config: cfg.LoggingConfig, max_iter: int, banner_messages curr_writer = config.local_writer.setup(banner_messages=banner_messages) EVENT_WRITERS.append(curr_writer) else: - logging.info("disabled local writer") + CONSOLE.log("disabled local writer") ## configure all the global buffer basic information GLOBAL_BUFFER["max_iter"] = max_iter @@ -388,10 +386,10 @@ class LocalWriter: valid_step = step > 0 and step % GLOBAL_BUFFER["steps_per_log"] == 0 if valid_step: if not self.has_printed and self.config.max_log_size: - logging.info( - "\x1b[33;20mPrinting max of %d lines. Set flag `--logging.local-writer.max-log-size=0` " - "to disable line wrapping.\x1b[0m", - self.config.max_log_size, + CONSOLE.log( + f"Printing max of {self.config.max_log_size} lines. " + "Set flag `--logging.local-writer.max-log-size=0` " + "to disable line wrapping." ) latest_map, new_key = self._consolidate_events() self._update_header(latest_map, new_key) @@ -433,6 +431,7 @@ class LocalWriter: self.past_mssgs[0] = mssg self.past_mssgs[1] = "-" * len(mssg) if full_log_cond or not self.has_printed: + print("testeing testing!") print(mssg) print("-" * len(mssg)) # self.has_printed = True diff --git a/nerfstudio/viewer/app/package.json b/nerfstudio/viewer/app/package.json index de8ca23e8..b0768523b 100644 --- a/nerfstudio/viewer/app/package.json +++ b/nerfstudio/viewer/app/package.json @@ -1,7 +1,7 @@ { "name": "viewer", "homepage": ".", - "version": "22-10-07-0", + "version": "22-10-13-0", "private": true, "dependencies": { "@emotion/react": "^11.10.4", diff --git a/nerfstudio/viewer/app/src/modules/Banner/Banner.jsx b/nerfstudio/viewer/app/src/modules/Banner/Banner.jsx index bf3c51d28..0ab082edd 100644 --- a/nerfstudio/viewer/app/src/modules/Banner/Banner.jsx +++ b/nerfstudio/viewer/app/src/modules/Banner/Banner.jsx @@ -17,22 +17,19 @@ function getParam(param_name) { return decodeURIComponent(params[1]); } -function getWebsocketEndpoint() { - const endpoint = getParam('websocket_port'); - return endpoint; -} - export default function Banner() { const dispatch = useDispatch(); let open_modal = true; - const websocket_port_from_argument = getWebsocketEndpoint(); - if (websocket_port_from_argument !== undefined) { + + // possibly set the websocket url + const websocket_url_from_argument = getParam('websocket_url'); + if (websocket_url_from_argument !== undefined) { open_modal = false; dispatch({ type: 'write', - path: 'websocketState/websocket_port', - data: websocket_port_from_argument, + path: 'websocketState/websocket_url', + data: websocket_url_from_argument, }); } diff --git a/nerfstudio/viewer/app/src/modules/LandingModal/LandingModal.jsx b/nerfstudio/viewer/app/src/modules/LandingModal/LandingModal.jsx index 6de67049a..645445680 100644 --- a/nerfstudio/viewer/app/src/modules/LandingModal/LandingModal.jsx +++ b/nerfstudio/viewer/app/src/modules/LandingModal/LandingModal.jsx @@ -13,7 +13,6 @@ import { } from '@mui/material'; import ContentCopyRoundedIcon from '@mui/icons-material/ContentCopyRounded'; -import { useSelector } from 'react-redux'; import WebSocketUrlField from '../WebSocketUrlField'; interface LandingModalProps { @@ -57,12 +56,9 @@ export default function LandingModel(props: LandingModalProps) { const handleOpen = () => setOpen(true); const handleClose = () => setOpen(false); - const text_intro = `To use our viewer, you need to connect to the WebSocket server where your training job is running. Enter your WebSocket Port below and click the link to connect to your training job.`; + const text_intro = `To use our viewer, you need to connect to the WebSocket server where your training job is running. Enter your WebSocket URL below and click the link to connect to your training job.`; - const websocket_port = useSelector( - (state) => state.websocketState.websocket_port, - ); - const command = `ssh -L ${websocket_port}:localhost:${websocket_port} USER@REMOTE.SERVER.IP`; + const command = `ssh -L :localhost: USER@REMOTE.SERVER.IP`; const platform = window.navigator.platform.toLowerCase(); let initial_tab = 0; @@ -180,9 +176,15 @@ export default function LandingModel(props: LandingModalProps) {

Camera Keyboard Controls

-

WASD - horizontal camera translation, QE - vertical camera translation
+
+ WASD - horizontal camera translation, QE - vertical camera + translation +
Arrow Keys - camera rotation
-
Spacebar - reset camera orbit point and view direction to origin
+
+ Spacebar - reset camera orbit point and view direction to + origin +

diff --git a/nerfstudio/viewer/app/src/modules/SidePanel/StatusPanel/StatusPanel.jsx b/nerfstudio/viewer/app/src/modules/SidePanel/StatusPanel/StatusPanel.jsx index 577618b6d..972b86518 100644 --- a/nerfstudio/viewer/app/src/modules/SidePanel/StatusPanel/StatusPanel.jsx +++ b/nerfstudio/viewer/app/src/modules/SidePanel/StatusPanel/StatusPanel.jsx @@ -108,6 +108,17 @@ export default function StatusPanel(props: StatusPanelProps) { > {cameras_button} +
Resolution: {eval_res} diff --git a/nerfstudio/viewer/app/src/modules/WebRtcWindow/WebRtcWindow.jsx b/nerfstudio/viewer/app/src/modules/WebRtcWindow/WebRtcWindow.jsx index a7012c499..900ee9591 100644 --- a/nerfstudio/viewer/app/src/modules/WebRtcWindow/WebRtcWindow.jsx +++ b/nerfstudio/viewer/app/src/modules/WebRtcWindow/WebRtcWindow.jsx @@ -122,8 +122,8 @@ export default function WebRtcWindow() { // send the offer console.log('[webrtc] sending offer'); const offer = pcRef.current.localDescription; - const cmd = 'offer'; - const path = ''; + const cmd = 'write'; + const path = 'webrtc/offer'; const data = { type: cmd, path, @@ -141,11 +141,13 @@ export default function WebRtcWindow() { websocket.addEventListener('message', (originalCmd) => { // set the remote description when the offer is received const cmd = msgpack.decode(new Uint8Array(originalCmd.data)); - if (cmd.type === 'answer') { + if (cmd.path === '/webrtc/answer') { console.log('[webrtc] received answer'); const answer = cmd.data; console.log(answer); - pcRef.current.setRemoteDescription(answer); + if (answer !== null) { + pcRef.current.setRemoteDescription(answer); + } } }); diff --git a/nerfstudio/viewer/app/src/modules/WebSocket/WebSocket.jsx b/nerfstudio/viewer/app/src/modules/WebSocket/WebSocket.jsx index 7b756ac88..0ac088ea9 100644 --- a/nerfstudio/viewer/app/src/modules/WebSocket/WebSocket.jsx +++ b/nerfstudio/viewer/app/src/modules/WebSocket/WebSocket.jsx @@ -15,15 +15,15 @@ export default function WebSocketContextFunction({ children }) { let ws = null; let socket = null; - // this code will rerender anytime the websocket changes now - const websocket_port = useSelector( - (state) => state.websocketState.websocket_port, + // this code will rerender anytime the websocket url changes + const websocket_url = useSelector( + (state) => state.websocketState.websocket_url, ); const connect = () => { - const url = `ws://localhost:${websocket_port}/`; - console.log(url); - socket = new WebSocket(url); + // of the form wss://ip_address:port + console.log(websocket_url); + socket = new WebSocket(websocket_url); socket.binaryType = 'arraybuffer'; socket.onopen = () => { dispatch({ @@ -58,7 +58,7 @@ export default function WebSocketContextFunction({ children }) { return () => { socket.close(); }; - }, [websocket_port]); + }, [websocket_url]); connect(); ws = { diff --git a/nerfstudio/viewer/app/src/modules/WebSocketUrlField.jsx b/nerfstudio/viewer/app/src/modules/WebSocketUrlField.jsx index c99754c1e..4dcec3376 100644 --- a/nerfstudio/viewer/app/src/modules/WebSocketUrlField.jsx +++ b/nerfstudio/viewer/app/src/modules/WebSocketUrlField.jsx @@ -4,15 +4,17 @@ import { TextField, Link } from '@mui/material'; import { useDispatch, useSelector } from 'react-redux'; export default function WebSocketUrlField() { - const websocket_port = useSelector( - (state) => state.websocketState.websocket_port, - ); const dispatch = useDispatch(); - const websocket_port_onchange = (event) => { + + // websocket url + const websocket_url = useSelector( + (state) => state.websocketState.websocket_url, + ); + const websocket_url_onchange = (event) => { const value = event.target.value; dispatch({ type: 'write', - path: 'websocketState/websocket_port', + path: 'websocketState/websocket_url', data: value, }); }; @@ -21,14 +23,14 @@ export default function WebSocketUrlField() {
- - viewer.nerf.studio?websocket_port={websocket_port} + + viewer.nerf.studio?websocket_url={websocket_url}
); diff --git a/nerfstudio/viewer/app/src/reducer.js b/nerfstudio/viewer/app/src/reducer.js index 956894b0e..ba1d8262b 100644 --- a/nerfstudio/viewer/app/src/reducer.js +++ b/nerfstudio/viewer/app/src/reducer.js @@ -3,8 +3,8 @@ import { split_path } from './utils'; const initialState = { // the websocket connection state websocketState: { - websocket_port: 7007, isConnected: false, + websocket_url: "ws://localhost:7007", }, // the webrtc connection state webrtcState: { diff --git a/nerfstudio/viewer/server/server.py b/nerfstudio/viewer/server/server.py index 6ff41ea87..c9c9e711d 100644 --- a/nerfstudio/viewer/server/server.py +++ b/nerfstudio/viewer/server/server.py @@ -17,8 +17,6 @@ import sys from typing import List, Optional, Tuple -import msgpack -import msgpack_numpy import tornado.gen import tornado.ioloop import tornado.web @@ -27,27 +25,11 @@ import tyro import umsgpack import zmq import zmq.eventloop.ioloop -from aiortc import RTCPeerConnection, RTCSessionDescription -from aiortc.rtcrtpsender import RTCRtpSender +from pyngrok import ngrok from zmq.eventloop.zmqstream import ZMQStream from nerfstudio.viewer.server.state.node import find_node, get_tree, walk from nerfstudio.viewer.server.state.state_node import StateNode -from nerfstudio.viewer.server.video_stream import SingleFrameStreamTrack - - -def force_codec(pc: RTCPeerConnection, sender: RTCRtpSender, forced_codec: str) -> None: - """Sets the codec preferences on a connection between sender and reciever - - Args: - pc: peer connection point - sender: sender that will send to connection point - forced_codec: codec to set - """ - kind = forced_codec.split("/")[0] - codecs = RTCRtpSender.getCapabilities(kind).codecs - transceiver = next(t for t in pc.getTransceivers() if t.sender == sender) - transceiver.setCodecPreferences([codec for codec in codecs if codec.mimeType == forced_codec]) class WebSocketHandler(tornado.websocket.WebSocketHandler): # pylint: disable=abstract-method @@ -82,7 +64,7 @@ class WebSocketHandler(tornado.websocket.WebSocketHandler): # pylint: disable=a if type_ == "write": # writes the data coming from the websocket find_node(self.bridge.state_tree, path).data = m["data"] - if m["path"] != "renderingState/camera": + if m["path"] != "renderingState/camera" and m["path"] != "webrtc/offer": # dispatch to the websockets if not a camera update! # TODO(ethan): handle the camera properly! # TODO: but don't update the current websocket @@ -94,29 +76,6 @@ class WebSocketHandler(tornado.websocket.WebSocketHandler): # pylint: disable=a # reads and returns the data data = find_node(self.bridge.state_tree, path).data self.write_message(data, binary=True) - elif type_ == "offer": - # returns the description to for WebRTC to the specific websocket connection - offer = RTCSessionDescription(m["data"]["sdp"], m["data"]["type"]) - - pc = RTCPeerConnection() - self.bridge.pcs.add(pc) - - video = SingleFrameStreamTrack() - self.bridge.video_tracks.add(video) - video_sender = pc.addTrack(video) - force_codec(pc, video_sender, "video/VP8") - - await pc.setRemoteDescription(offer) - answer = await pc.createAnswer() - await pc.setLocalDescription(answer) - - cmd_data = { - "type": "answer", - "path": "", - "data": {"sdp": pc.localDescription.sdp, "type": pc.localDescription.type}, - } - data = umsgpack.packb(cmd_data) - self.write_message(data, binary=True) else: cmd_data = { "type": "error", @@ -136,7 +95,6 @@ class ZMQWebSocketBridge: Args: 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. """ @@ -147,11 +105,9 @@ class ZMQWebSocketBridge: self.websocket_pool = set() self.app = self.make_app() self.ioloop = tornado.ioloop.IOLoop.current() - self.pcs = set() - self.video_tracks = set() # zmq - zmq_url = f"tcp://127.0.0.1:{self.zmq_port:d}" + zmq_url = f"tcp://0.0.0.0:{self.zmq_port:d}" self.zmq_socket, self.zmq_stream, self.zmq_url = self.setup_zmq(zmq_url) # websocket @@ -195,13 +151,6 @@ class ZMQWebSocketBridge: # TODO(ethan): handle the "data" key... read_data = find_node(self.state_tree, path).data self.zmq_socket.send(umsgpack.packb(read_data)) - elif type_ == "set_image": - image = msgpack.unpackb( - data, object_hook=msgpack_numpy.decode, use_list=False, max_bin_len=50000000, raw=False - ) - for video_track in self.video_tracks: - video_track.put_frame(image) - self.zmq_socket.send(umsgpack.packb(b"ok")) else: self.zmq_socket.send(umsgpack.packb(b"error: unknown command")) @@ -249,13 +198,22 @@ class ZMQWebSocketBridge: self.ioloop.start() -def run_viewer_bridge_server(zmq_port: int = 6000, websocket_port: int = 7007): +def run_viewer_bridge_server(zmq_port: int = 6000, websocket_port: int = 7007, use_ngrok: bool = False): """Run the viewer bridge server. Args: zmq_port: port to use for zmq websocket_port: port to use for websocket + host: host to connect to """ + + # whether to launch pyngrok or not + if use_ngrok: + # Open a HTTP tunnel on the default port 80 + # .ngrok.io" -> "http://localhost:80"> + http_tunnel = ngrok.connect(addr=str(zmq_port), proto="tcp") + print(http_tunnel) + bridge = ZMQWebSocketBridge(zmq_port=zmq_port, websocket_port=websocket_port) print(bridge) try: @@ -264,11 +222,11 @@ def run_viewer_bridge_server(zmq_port: int = 6000, websocket_port: int = 7007): pass -def main(): +def entrypoint(): """The main entrypoint.""" tyro.extras.set_accent_color("bright_yellow") tyro.cli(run_viewer_bridge_server) if __name__ == "__main__": - main() + entrypoint() diff --git a/nerfstudio/viewer/server/utils.py b/nerfstudio/viewer/server/utils.py index 159665ed7..92ccc0ef7 100644 --- a/nerfstudio/viewer/server/utils.py +++ b/nerfstudio/viewer/server/utils.py @@ -21,6 +21,8 @@ from typing import Any, Callable, Dict, List, Optional, Tuple import numpy as np import torch import zmq +from aiortc import RTCPeerConnection +from aiortc.rtcrtpsender import RTCRtpSender def get_chunks( @@ -110,3 +112,17 @@ def find_available_port(func: Callable, default_port: int, max_attempts: int = 1 raise ( Exception(f"Could not find an available port in the range: [{default_port:d}, {max_attempts + default_port:d})") ) + + +def force_codec(pc: RTCPeerConnection, sender: RTCRtpSender, forced_codec: str) -> None: + """Sets the codec preferences on a connection between sender and reciever + + Args: + pc: peer connection point + sender: sender that will send to connection point + forced_codec: codec to set + """ + kind = forced_codec.split("/")[0] + codecs = RTCRtpSender.getCapabilities(kind).codecs + transceiver = next(t for t in pc.getTransceivers() if t.sender == sender) + transceiver.setCodecPreferences([codec for codec in codecs if codec.mimeType == forced_codec]) diff --git a/nerfstudio/viewer/server/video_stream.py b/nerfstudio/viewer/server/video_stream.py index a0c3e7cef..c50168055 100644 --- a/nerfstudio/viewer/server/video_stream.py +++ b/nerfstudio/viewer/server/video_stream.py @@ -14,66 +14,11 @@ """Video Stream objects for WebRTC""" -import math - -import cv2 import numpy as np from aiortc import VideoStreamTrack from av import VideoFrame -class FlagVideoStreamTrack(VideoStreamTrack): - """ - A video track that returns an animated flag. For debugging purposes :') - """ - - def __init__(self): - super().__init__() # don't forget this! - self.counter = 0 - height, width = 480, 640 - - # generate flag - data_bgr = np.hstack( - [ - self._create_rectangle(width=213, height=480, color=(255, 0, 0)), # blue - self._create_rectangle(width=214, height=480, color=(255, 255, 255)), # white - self._create_rectangle(width=213, height=480, color=(0, 0, 255)), # red - ] - ) - - # shrink and center it - M = np.float32([[0.5, 0, width / 4], [0, 0.5, height / 4]]) - data_bgr = cv2.warpAffine(data_bgr, M, (width, height)) - - # compute animation - omega = 2 * math.pi / height - id_x = np.tile(np.array(range(width), dtype=np.float32), (height, 1)) - id_y = np.tile(np.array(range(height), dtype=np.float32), (width, 1)).transpose() - - self.frames = [] - for k in range(30): - phase = 2 * k * math.pi / 30 - map_x = id_x + 10 * np.cos(omega * id_x + phase) - map_y = id_y + 10 * np.sin(omega * id_x + phase) - frame = cv2.remap(data_bgr, map_x, map_y, cv2.INTER_LINEAR) - self.frames.append(VideoFrame.from_ndarray(frame, format="bgr24")) - - async def recv(self): - pts, time_base = await self.next_timestamp() - - frame = self.frames[self.counter % 30] - frame.pts = pts - frame.time_base = time_base - self.counter += 1 - return frame - - @staticmethod - def _create_rectangle(width, height, color): - data_bgr = np.zeros((height, width, 3), np.uint8) - data_bgr[:, :] = color - return data_bgr - - class SingleFrameStreamTrack(VideoStreamTrack): """Single Frame stream class: pushes single frames to a stream""" diff --git a/nerfstudio/viewer/server/viewer_utils.py b/nerfstudio/viewer/server/viewer_utils.py index b246c0612..f3bf7f7d0 100644 --- a/nerfstudio/viewer/server/viewer_utils.py +++ b/nerfstudio/viewer/server/viewer_utils.py @@ -16,16 +16,20 @@ """ from __future__ import annotations +import asyncio import enum import os import sys import threading import time +import warnings from pathlib import Path from typing import Any, Dict, Optional, Tuple import numpy as np import torch +from aiortc import RTCPeerConnection, RTCSessionDescription +from cryptography.utils import CryptographyDeprecationWarning from rich.console import Console from nerfstudio.cameras.cameras import Cameras @@ -39,10 +43,16 @@ from nerfstudio.utils.io import load_from_json, write_to_json from nerfstudio.utils.misc import get_dict_to_torch from nerfstudio.utils.writer import GLOBAL_BUFFER, EventName, TimeWriter from nerfstudio.viewer.server.subprocess import run_viewer_bridge_server_as_subprocess -from nerfstudio.viewer.server.utils import get_intrinsics_matrix_and_camera_to_world_h +from nerfstudio.viewer.server.utils import ( + force_codec, + get_intrinsics_matrix_and_camera_to_world_h, +) +from nerfstudio.viewer.server.video_stream import SingleFrameStreamTrack from nerfstudio.viewer.server.visualizer import Viewer -console = Console(width=120) +warnings.filterwarnings("ignore", category=CryptographyDeprecationWarning) + +CONSOLE = Console(width=120) @check_main_thread @@ -210,24 +220,25 @@ class ViewerState: self.log_filename = log_filename if self.config.launch_bridge_server: # start the viewer bridge server - websocket_port = self.config.websocket_port + assert self.config.websocket_port is not None self.log_filename.parent.mkdir(exist_ok=True) zmq_port = run_viewer_bridge_server_as_subprocess( - websocket_port, zmq_port=self.config.zmq_port, log_filename=str(self.log_filename) + self.config.websocket_port, zmq_port=self.config.zmq_port, log_filename=str(self.log_filename) ) # TODO(ethan): log the output of the viewer bridge server in a file where the training logs go - console.line() + 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.nerf.studio/versions/{version}/?websocket_port={websocket_port}" - console.rule(characters="=") - console.print(f"[Public] Open the viewer at {self.viewer_url}") - console.rule(characters="=") - console.line() + websocket_url = f"ws://localhost:{self.config.websocket_port}" + self.viewer_url = f"https://viewer.nerf.studio/versions/{version}/?websocket_url={websocket_url}" + CONSOLE.rule(characters="=") + CONSOLE.print(f"[Public] Open the viewer at {self.viewer_url}") + CONSOLE.rule(characters="=") + CONSOLE.line() self.vis = Viewer(zmq_port=zmq_port) else: assert self.config.zmq_port is not None - self.vis = Viewer(zmq_port=self.config.zmq_port) + self.vis = Viewer(zmq_port=self.config.zmq_port, ip_address=self.config.ip_address) # viewer specific variables self.prev_camera_matrix = None @@ -246,6 +257,12 @@ class ViewerState: self.output_list = None + # webrtc + self.pcs = set() + self.video_tracks = set() + self.webrtc_thread = None + self.kill_webrtc_signal = False + def init_scene(self, dataset: InputDataset, start_train=True) -> None: """Draw some images and the scene aabb in the viewer. @@ -296,6 +313,26 @@ class ViewerState: write_to_json(Path(camera_path_filename), camera_path) self.vis["camera_path_payload"].delete() + def _check_webrtc_offer(self): + """Check if there is a webrtc offer to respond to.""" + data = self.vis["webrtc/offer"].read() + if data: + if self.webrtc_thread and self.webrtc_thread.is_alive(): + # kill the previous thread if the webpage refreshes + self.kill_webrtc_signal = True + return + + def loop_in_thread(loop): + asyncio.set_event_loop(loop) + loop.run_until_complete(self.send_webrtc_answer(data)) + + loop = asyncio.get_event_loop() + self.webrtc_thread = threading.Thread(target=loop_in_thread, args=(loop,)) + self.webrtc_thread.daemon = True + self.webrtc_thread.start() + # remove the offer from the state tree + self.vis["webrtc/offer"].delete() + def update_scene(self, trainer, step: int, graph: Model, num_rays_per_batch: int) -> None: """updates the scene based on the graph weights @@ -308,6 +345,7 @@ class ViewerState: self.step = step self._check_camera_path_payload(trainer, step) + self._check_webrtc_offer() camera_object = self._get_camera_object() if camera_object is None: @@ -352,6 +390,7 @@ class ViewerState: camera_object = self._get_camera_object() is_training = self.vis["renderingState/isTraining"].read() self._check_camera_path_payload(trainer, step) + self._check_webrtc_offer() run_loop = not is_training local_step += 1 @@ -431,6 +470,39 @@ class ViewerState: raise NotImplementedError + async def send_webrtc_answer(self, data): + """Setup the webrtc connection.""" + + # returns the description to for WebRTC to the specific websocket connection + offer = RTCSessionDescription(data["sdp"], data["type"]) + + pc = RTCPeerConnection() + self.pcs.add(pc) + + video = SingleFrameStreamTrack() + self.video_tracks.add(video) + video_sender = pc.addTrack(video) + force_codec(pc, video_sender, "video/VP8") + + await pc.setRemoteDescription(offer) + answer = await pc.createAnswer() + await pc.setLocalDescription(answer) + + self.vis["webrtc/answer"].write({"sdp": pc.localDescription.sdp, "type": pc.localDescription.type}) + self.vis["webrtc/answer"].delete() + + # continually exchange media + while True: + await asyncio.sleep(1) + if self.kill_webrtc_signal: + self.kill_webrtc_signal = False + return + + def set_image(self, image): + """Write the image over webrtc.""" + for video_track in self.video_tracks: + video_track.put_frame(image) + def _send_output_to_viewer(self, outputs: Dict[str, Any], stuff_colors: torch.Tensor = None, eps=1e-6): """Chooses the correct output and sends it to the viewer @@ -465,7 +537,7 @@ class ViewerState: self.vis["renderingState/colormap_options"].write(colormap_options) selected_output = (self._apply_colormap(outputs, stuff_colors) * 255).type(torch.uint8) image = selected_output.cpu().numpy() - self.vis.set_image(image) + self.set_image(image) def _update_viewer_stats(self, render_time: float, num_rays: int, image_height: int, image_width: int) -> None: """Function that calculates and populates all the rendering statistics accordingly diff --git a/nerfstudio/viewer/server/visualizer.py b/nerfstudio/viewer/server/visualizer.py index e7d26b63d..2ec367b58 100644 --- a/nerfstudio/viewer/server/visualizer.py +++ b/nerfstudio/viewer/server/visualizer.py @@ -15,28 +15,33 @@ """Code to connect and send commands to the viewer. """ -import logging import sys from threading import Thread from typing import Dict, Optional, Union -import msgpack -import msgpack_numpy import umsgpack import zmq +from rich.console import Console from nerfstudio.viewer.server.path import Path +CONSOLE = Console(width=120) + class ViewerWindow: - """The viewer window has the ZMQ connection to the viewer bridge server.""" + """The viewer window has the ZMQ connection to the viewer bridge server. + + Args: + zmq_port: Where to connect with ZMQ. + ip_address: The ip address of the bridge server. + """ context = zmq.Context() - def __init__(self, zmq_port): + def __init__(self, zmq_port, ip_address="0.0.0.0"): self.zmq_port = zmq_port self.client = self.context.socket(zmq.REQ) - zmq_url = f"tcp://127.0.0.1:{self.zmq_port}" + zmq_url = f"tcp://{ip_address}:{self.zmq_port}" self.client.connect(zmq_url) self.assert_connected() @@ -79,7 +84,7 @@ class ViewerWindow: t.start() t.join(timeout_in_sec) except Exception as je: - logging.info("Error starting thread") + CONSOLE.log("Error starting thread") raise je ret = res[0] if isinstance(ret, BaseException): @@ -93,12 +98,12 @@ class ViewerWindow: timeout_in_sec (int): The maximum time to wait for the connection to be established. """ try: - logging.info("Sending ping to the viewer Bridge Server...") + CONSOLE.print("Sending ping to the viewer Bridge Server...") _ = self.timeout_ping(timeout_in_sec) - logging.info("Successfully connected.") + CONSOLE.print("Successfully connected.") except Exception as e: # pylint: disable=broad-except - logging.info(e) + CONSOLE.log(e) sys.exit() @@ -108,13 +113,16 @@ class Viewer: Args: zmq_port: Where to connect with ZMQ. window: An already existing ViewerWindow. + ip_address: The ip address of the bridge server. """ - def __init__(self, zmq_port: Optional[int] = None, window: Optional[ViewerWindow] = None): + def __init__( + self, zmq_port: Optional[int] = None, window: Optional[ViewerWindow] = None, ip_address: str = "0.0.0.0" + ): if zmq_port is None and window is None: raise ValueError("Must specify either zmq_port or window.") if window is None: - self.window = ViewerWindow(zmq_port=zmq_port) + self.window = ViewerWindow(zmq_port=zmq_port, ip_address=ip_address) else: self.window = window self.path = Path(()) @@ -145,17 +153,3 @@ class Viewer: def delete(self): """Delete data.""" return self.write(data=None) - - def set_image(self, image): - """Sends an image to the viewer with WebRTC.""" - type_ = "set_image" - path = self.path.lower() - data = msgpack.packb(image, default=msgpack_numpy.encode, use_bin_type=True) - self.window.client.send_multipart( - [ - type_.encode("utf-8"), - path.encode("utf-8"), - data, - ] - ) - return self.window.client.recv() diff --git a/pyproject.toml b/pyproject.toml index ce024638d..ac6863bf9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,6 +34,7 @@ dependencies = [ "opencv-python==4.6.0.66", "plotly==5.7.0", "protobuf==3.20.0", + "pyngrok==5.1.0", "python-socketio==5.7.1", "requests", "rich==12.5.1", @@ -83,6 +84,7 @@ ns-train = "scripts.train:entrypoint" ns-eval = "scripts.eval:entrypoint" ns-render = "scripts.render:entrypoint" ns-dev-test = "scripts.github.run_actions:entrypoint" +ns-bridge-server = "nerfstudio.viewer.server.server:entrypoint" [options] # equivalent to using --extra-index-url with pip, which is needed for specifying the CUDA version torch and torchvision diff --git a/scripts/completions/install.py b/scripts/completions/install.py index 2ec668f08..3d413aa16 100755 --- a/scripts/completions/install.py +++ b/scripts/completions/install.py @@ -137,6 +137,9 @@ def _update_rc( [ "", "# Source nerfstudio autocompletions.", + "# From https://unix.stackexchange.com/questions/339954/zsh-command-not-found-compinstall-compinit-compdef", + "autoload -Uz compinit", + "compinit", f"source {completions_dir / 'setup'}.{shell}", ] ) diff --git a/scripts/docs/add_nb_tags.py b/scripts/docs/add_nb_tags.py index 73edc5a84..37a5900f8 100755 --- a/scripts/docs/add_nb_tags.py +++ b/scripts/docs/add_nb_tags.py @@ -7,7 +7,7 @@ import nbformat as nbf import tyro from rich.console import Console -console = Console(width=120) +CONSOLE = Console(width=120) def main(check: bool = False): @@ -47,7 +47,7 @@ def main(check: bool = False): found_tags.append(val) if len(found_keys) > 1: - console.print(f"[bold yellow]Found multiple tags {found_keys} for {ipath}") + CONSOLE.print(f"[bold yellow]Found multiple tags {found_keys} for {ipath}") sys.exit(1) if len(cell_tags) != len(found_tags): @@ -59,7 +59,7 @@ def main(check: bool = False): cell["metadata"]["tags"] = found_tags if incorrect_metadata: if check: - console.print( + CONSOLE.print( f"[bold yellow]{ipath} has incorrect metadata. Call `python scripts.docs.add_nb_tags.py` to add it." ) any_missing = True @@ -68,7 +68,7 @@ def main(check: bool = False): nbf.write(ntbk, ipath) if not any_missing: - console.print("[green]All notebooks have correct metadata.") + CONSOLE.print("[green]All notebooks have correct metadata.") if check and any_missing: sys.exit(1) diff --git a/scripts/docs/build_docs.py b/scripts/docs/build_docs.py index 641bce41d..cb31e54b3 100755 --- a/scripts/docs/build_docs.py +++ b/scripts/docs/build_docs.py @@ -7,7 +7,7 @@ import tyro from rich.console import Console from rich.style import Style -console = Console(width=120) +CONSOLE = Console(width=120) LOCAL_TESTS = ["Run license checks", "Run Black", "Python Pylint", "Test with pytest"] @@ -20,7 +20,7 @@ def run_command(command: str) -> None: """ ret_code = subprocess.call(command, shell=True) if ret_code != 0: - console.print(f"[bold red]Error: `{command}` failed. Exiting...") + CONSOLE.print(f"[bold red]Error: `{command}` failed. Exiting...") sys.exit(1) @@ -31,20 +31,20 @@ def main(clean_cache: bool = False): clean_cache: whether to clean the cache before building docs. """ - console.print("[green]Adding notebook documentation metadata") + CONSOLE.print("[green]Adding notebook documentation metadata") run_command("python scripts/docs/add_nb_tags.py") # Add checks for building documentation - console.print("[green]Building Documentation") + CONSOLE.print("[green]Building Documentation") if clean_cache: run_command("cd docs/; make clean; make html SPHINXOPTS='-W;'") else: run_command("cd docs/; make html SPHINXOPTS='-W;'") - console.line() - console.rule(characters="=", style=Style(color="green")) - console.print("[bold green]Done") - console.rule(characters="=", style=Style(color="green")) + CONSOLE.line() + CONSOLE.rule(characters="=", style=Style(color="green")) + CONSOLE.print("[bold green]Done") + CONSOLE.rule(characters="=", style=Style(color="green")) if __name__ == "__main__": diff --git a/scripts/downloads/download_data.py b/scripts/downloads/download_data.py index fdb0eb9d4..d75e4a1f1 100644 --- a/scripts/downloads/download_data.py +++ b/scripts/downloads/download_data.py @@ -10,7 +10,7 @@ import gdown import tyro from rich.console import Console -console = Console(width=120) +CONSOLE = Console(width=120) def download_blender(save_dir: Path): @@ -131,19 +131,19 @@ def main( if dataset == "blender": if capture_name is not None: - console.print("Capture name is ignored when downloading from the blender dataset.") + CONSOLE.print("Capture name is ignored when downloading from the blender dataset.") download_blender(save_dir) if dataset == "friends": if capture_name is not None: - console.print("Capture name is ignored when downloading from the blender dataset.") + CONSOLE.print("Capture name is ignored when downloading from the blender dataset.") download_friends(save_dir) if dataset == "nerfstudio": if capture_name is None: capture_names = sorted(nerfstudio_file_ids.keys()) - console.rule("[bold red]Error", style="bold red") - console.print("[bold yellow]You must pass in --capture-name when downloading from the nerfstudio dataset.") - console.print("Use one of the following:") - console.print(f"\t {capture_names}") + CONSOLE.rule("[bold red]Error", style="bold red") + CONSOLE.print("[bold yellow]You must pass in --capture-name when downloading from the nerfstudio dataset.") + CONSOLE.print("Use one of the following:") + CONSOLE.print(f"\t {capture_names}") sys.exit() download_nerfstudio(save_dir, capture_name) if dataset == "record3d": diff --git a/scripts/eval.py b/scripts/eval.py index 6f5833665..b34a04373 100755 --- a/scripts/eval.py +++ b/scripts/eval.py @@ -5,7 +5,6 @@ eval.py from __future__ import annotations import json -import logging from dataclasses import dataclass from pathlib import Path @@ -14,9 +13,7 @@ from rich.console import Console from nerfstudio.utils.eval_utils import eval_setup -console = Console(width=120) - -logging.basicConfig(format="[%(filename)s:%(lineno)d] %(message)s", level=logging.INFO) +CONSOLE = Console(width=120) @dataclass @@ -43,7 +40,7 @@ class ComputePSNR: "results": metrics_dict, } self.output_path.write_text(json.dumps(benchmark_info, indent=2), "utf8") - console.print(f"Saved results to: {self.output_path}") + CONSOLE.print(f"Saved results to: {self.output_path}") def entrypoint(): diff --git a/scripts/github/run_actions.py b/scripts/github/run_actions.py index 299a83624..19179725c 100755 --- a/scripts/github/run_actions.py +++ b/scripts/github/run_actions.py @@ -8,7 +8,7 @@ import yaml from rich.console import Console from rich.style import Style -console = Console(width=120) +CONSOLE = Console(width=120) LOCAL_TESTS = ["Run license checks", "Run isort", "Run Black", "Python Pylint", "Test with pytest"] @@ -22,7 +22,7 @@ def run_command(command: str, continue_on_fail: bool = False) -> bool: """ ret_code = subprocess.call(command, shell=True) if ret_code != 0: - console.print(f"[bold red]Error: `{command}` failed.") + CONSOLE.print(f"[bold red]Error: `{command}` failed.") if not continue_on_fail: sys.exit(1) return ret_code == 0 @@ -47,31 +47,31 @@ def run_github_actions_file(filename: str, continue_on_fail: bool = False): compressed = compressed.replace("--check", "") curr_command = f"{compressed}" - console.line() - console.rule(f"[bold green]Running: {curr_command}") + CONSOLE.line() + CONSOLE.rule(f"[bold green]Running: {curr_command}") success = success and run_command(curr_command, continue_on_fail=continue_on_fail) else: skip_name = step["name"] if "name" in step else step["uses"] - console.print(f"Skipping {skip_name}") + CONSOLE.print(f"Skipping {skip_name}") # Add checks for building documentation - console.line() - console.rule("[bold green]Adding notebook documentation metadata") + CONSOLE.line() + CONSOLE.rule("[bold green]Adding notebook documentation metadata") success = success and run_command("python scripts/docs/add_nb_tags.py", continue_on_fail=continue_on_fail) - console.line() - console.rule("[bold green]Building Documentation") + CONSOLE.line() + CONSOLE.rule("[bold green]Building Documentation") success = success and run_command("cd docs/; make html SPHINXOPTS='-W;'", continue_on_fail=continue_on_fail) if success: - console.line() - console.rule(characters="=") - console.print("[bold green]:TADA: :TADA: :TADA: ALL CHECKS PASSED :TADA: :TADA: :TADA:", justify="center") - console.rule(characters="=") + CONSOLE.line() + CONSOLE.rule(characters="=") + CONSOLE.print("[bold green]:TADA: :TADA: :TADA: ALL CHECKS PASSED :TADA: :TADA: :TADA:", justify="center") + CONSOLE.rule(characters="=") else: - console.line() - console.rule(characters="=", style=Style(color="red")) - console.print("[bold red]:skull: :skull: :skull: ERRORS FOUND :skull: :skull: :skull:", justify="center") - console.rule(characters="=", style=Style(color="red")) + CONSOLE.line() + CONSOLE.rule(characters="=", style=Style(color="red")) + CONSOLE.print("[bold red]:skull: :skull: :skull: ERRORS FOUND :skull: :skull: :skull:", justify="center") + CONSOLE.rule(characters="=", style=Style(color="red")) def run_code_checks(continue_on_fail: bool = False): diff --git a/scripts/render.py b/scripts/render.py index 89b92c8a1..5ada54af4 100644 --- a/scripts/render.py +++ b/scripts/render.py @@ -6,7 +6,6 @@ from __future__ import annotations import dataclasses import json -import logging import sys from pathlib import Path from typing import Literal, Optional @@ -32,9 +31,7 @@ from nerfstudio.utils import install_checks from nerfstudio.utils.eval_utils import eval_setup from nerfstudio.utils.rich_utils import ItersPerSecColumn -console = Console(width=120) - -logging.basicConfig(format="[%(filename)s:%(lineno)d] %(message)s", level=logging.INFO) +CONSOLE = Console(width=120) def _render_trajectory_video( @@ -55,7 +52,7 @@ def _render_trajectory_video( rendered_resolution_scaling_factor: Scaling factor to apply to the camera image resolution. seconds: Number for the output video. """ - console.print("[bold green]Creating trajectory video") + CONSOLE.print("[bold green]Creating trajectory video") images = [] cameras.rescale_output_resolution(rendered_resolution_scaling_factor) @@ -72,9 +69,9 @@ def _render_trajectory_video( with torch.no_grad(): outputs = pipeline.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle) if rendered_output_name not in outputs: - console.rule("Error", style="red") - console.print(f"Could not find {rendered_output_name} in the model outputs", justify="center") - console.print(f"Please set --rendered_output_name to one of: {outputs.keys()}", justify="center") + CONSOLE.rule("Error", style="red") + CONSOLE.print(f"Could not find {rendered_output_name} in the model outputs", justify="center") + CONSOLE.print(f"Please set --rendered_output_name to one of: {outputs.keys()}", justify="center") sys.exit(1) image = outputs[rendered_output_name].cpu().numpy() images.append(image) @@ -82,10 +79,10 @@ def _render_trajectory_video( fps = len(images) / seconds # make the folder if it doesn't exist output_filename.parent.mkdir(parents=True, exist_ok=True) - with console.status("[yellow]Saving video", spinner="bouncingBall"): + with CONSOLE.status("[yellow]Saving video", spinner="bouncingBall"): media.write_video(output_filename, images, fps=fps) - console.rule("[green] :tada: :tada: :tada: Success :tada: :tada: :tada:") - console.print(f"[green]Saved video to {output_filename}", justify="center") + CONSOLE.rule("[green] :tada: :tada: :tada: Success :tada: :tada: :tada:") + CONSOLE.print(f"[green]Saved video to {output_filename}", justify="center") @dataclasses.dataclass diff --git a/scripts/train.py b/scripts/train.py index 47afd26b1..86dd77d68 100755 --- a/scripts/train.py +++ b/scripts/train.py @@ -31,7 +31,6 @@ subcommand: from __future__ import annotations -import logging import random import socket import traceback @@ -44,6 +43,7 @@ import torch.distributed as dist import torch.multiprocessing as mp import tyro import yaml +from rich.console import Console from nerfstudio.configs import base_config as cfg from nerfstudio.configs.config_utils import convert_markup_to_ansi @@ -51,8 +51,7 @@ from nerfstudio.configs.method_configs import AnnotatedBaseConfigUnion from nerfstudio.engine.trainer import train_loop from nerfstudio.utils import comms, profiler -logging.basicConfig(format="[%(filename)s:%(lineno)d] %(message)s", level=logging.DEBUG) - +CONSOLE = Console(width=120) DEFAULT_TIMEOUT = timedelta(minutes=30) # speedup for when input size to model doesn't change (much) @@ -117,7 +116,7 @@ def _distributed_worker( timeout=timeout, ) except Exception as e: - logging.error("Process group URL: %s", dist_url) + CONSOLE.log(f"Process group URL: {dist_url}") raise e assert comms._LOCAL_PROCESS_GROUP is None # pylint: disable=protected-access @@ -181,7 +180,7 @@ def launch( port = _find_free_port() dist_url = f"tcp://127.0.0.1:{port}" if num_machines > 1 and dist_url.startswith("file://"): - logging.warning("file:// is not a reliable init_method in multi-machine jobs. Prefer tcp://") + CONSOLE.log("file:// is not a reliable init_method in multi-machine jobs. Prefer tcp://") process_context = mp.spawn( _distributed_worker, @@ -202,10 +201,10 @@ def launch( except KeyboardInterrupt: for i, process in enumerate(process_context.processes): if process.is_alive(): - logging.info("Terminating process %s", str(i)) + CONSOLE.log("Terminating process %s", str(i)) process.terminate() process.join() - logging.info("Process %s finished", str(i)) + CONSOLE.log("Process %s finished", str(i)) finally: profiler.flush_profiler(config.logging) @@ -215,11 +214,11 @@ def main(config: cfg.Config) -> None: config.set_timestamp() if config.data: - logging.info("Using --data alias for --data.pipeline.datamanager.dataparser.data") + CONSOLE.log("Using --data alias for --data.pipeline.datamanager.dataparser.data") config.pipeline.datamanager.dataparser.data = config.data if config.trainer.load_config: - logging.info(f"Loading pre-set config from: {config.trainer.load_config}") + CONSOLE.log(f"Loading pre-set config from: {config.trainer.load_config}") config = yaml.load(config.trainer.load_config.read_text(), Loader=yaml.Loader) # print and save config diff --git a/scripts/viewer/run_bridge_server.py b/scripts/viewer/run_bridge_server.py deleted file mode 100644 index 7aa27a348..000000000 --- a/scripts/viewer/run_bridge_server.py +++ /dev/null @@ -1,7 +0,0 @@ -""" -Runs the viewer bridge server. -""" -from nerfstudio.viewer.server.server import main - -if __name__ == "__main__": - main() diff --git a/scripts/viewer/view_dataset.py b/scripts/viewer/view_dataset.py index 8be8efb8c..29317515f 100644 --- a/scripts/viewer/view_dataset.py +++ b/scripts/viewer/view_dataset.py @@ -3,7 +3,9 @@ view_dataset.py """ -import logging +from rich.console import Console + +CONSOLE = Console(width=120) import time from datetime import timedelta from pathlib import Path @@ -11,13 +13,11 @@ from pathlib import Path import torch import tyro -from nerfstudio.configs import model_configs as cfg +from nerfstudio.configs import method_configs as cfg from nerfstudio.data.datamanagers import AnnotatedDataParserUnion from nerfstudio.data.utils.datasets import InputDataset from nerfstudio.viewer.server import viewer_utils -logging.basicConfig(format="[%(filename)s:%(lineno)d] %(message)s", level=logging.INFO) - DEFAULT_TIMEOUT = timedelta(minutes=30) # speedup for when input size to model doesn't change (much) @@ -36,7 +36,7 @@ def main( ) dataset = InputDataset(dataparser.setup().get_dataset_inputs(split="train")) viewer_state.init_scene(dataset=dataset, start_train=False) - logging.info("Please refresh and load page at: %s", viewer_state.viewer_url) + CONSOLE.log("Please refresh and load page at: %s", viewer_state.viewer_url) time.sleep(30) # allowing time to refresh page