From 56661b5d9aa8adfec9cad60bce53036cb0ceca43 Mon Sep 17 00:00:00 2001 From: Ethan Weber Date: Wed, 6 Jul 2022 15:47:08 -0700 Subject: [PATCH] Viewer Dynamic Resolution Scaling (DRS) (#87) * updated three.js version. only send camera information when it moves * changing code back to updating camera constantly because packets seem to be lost * interrupting function during rendering * small linting fixes * added logic for interrupting rendering + separted out visualizer logic from trainer Co-authored-by: evonneng --- configs/graph_default.yaml | 4 +- configs/graph_instant_ngp.yaml | 5 +- notebooks/visualize_viewer.ipynb | 158 ++++++---------- pyrad/engine/trainer.py | 85 +-------- pyrad/graphs/base.py | 40 ++-- pyrad/utils/config.py | 13 +- pyrad/viewer/app/package.json | 2 +- pyrad/viewer/app/src/SceneNode.js | 229 ++++++++++++----------- pyrad/viewer/app/src/Viewer.js | 19 ++ pyrad/viewer/app/yarn.lock | 8 +- pyrad/viewer/backend/run_zmq_server.py | 3 +- pyrad/viewer/backend/vis_utils.py | 247 ++++++++++++++++++------- 12 files changed, 418 insertions(+), 395 deletions(-) diff --git a/configs/graph_default.yaml b/configs/graph_default.yaml index 7e2bfa158..0d28c5d8e 100644 --- a/configs/graph_default.yaml +++ b/configs/graph_default.yaml @@ -102,6 +102,6 @@ optimizers: viewer: enable: False zmq_url: tcp://127.0.0.1:6000 - render_image_height: 100 - steps_per_render_image: 100 + min_render_image_height: 100 + max_render_image_height: 1000 num_rays_per_chunk: 4096 diff --git a/configs/graph_instant_ngp.yaml b/configs/graph_instant_ngp.yaml index 63bd04190..f42464434 100644 --- a/configs/graph_instant_ngp.yaml +++ b/configs/graph_instant_ngp.yaml @@ -25,7 +25,4 @@ optimizers: scheduler: null viewer: - enable: True - render_image_height: 200 - steps_per_render_image: 5 - num_rays_per_chunk: 16384 \ No newline at end of file + enable: True \ No newline at end of file diff --git a/notebooks/visualize_viewer.ipynb b/notebooks/visualize_viewer.ipynb index f1f71238b..873055d6b 100644 --- a/notebooks/visualize_viewer.ipynb +++ b/notebooks/visualize_viewer.ipynb @@ -10,7 +10,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 1, "id": "8501e01d-f932-4a75-8259-2def67a546d7", "metadata": {}, "outputs": [], @@ -23,7 +23,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 2, "id": "b15dfdc1-5054-40bc-9631-4c7426eecfd9", "metadata": {}, "outputs": [], @@ -67,42 +67,18 @@ }, { "cell_type": "code", - "execution_count": null, - "id": "032e4c9e-25be-4cc0-9ac8-a5be723b775a", - "metadata": {}, - "outputs": [], - "source": [ - "# establish the connection to the tcp backend, to talk to the visualizer\n", - "# zmq_url should match the output of `python run_zmq_server.py`\n", - "vis = vis_utils.get_vis(zmq_url=\"tcp://127.0.0.1:6000\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "df82e083-e261-43e0-a1dc-05c6a59d16e1", - "metadata": {}, - "outputs": [], - "source": [ - "vis.delete()\n", - "\n", - "# draws a red box in the scene\n", - "vis_utils.show_box_test(vis)" - ] - }, - { - "cell_type": "code", - "execution_count": null, + "execution_count": 3, "id": "7fe7478d", "metadata": {}, "outputs": [], "source": [ "with initialize(version_base=None, config_path=\"../configs\"):\n", " config = compose(\n", - " config_name=\"graph_default.yaml\",\n", + " config_name=\"graph_instant_ngp.yaml\",\n", " overrides=[\n", - " \"trainer.resume_train.load_dir=/shared/ethanweber/pyrad/outputs/blender_lego/vanilla_nerf/2022-07-04_221907/checkpoints\",\n", - " \"trainer.resume_train.load_step=4000\",\n", + " \"trainer.resume_train.load_dir=/shared/ethanweber/pyrad/outputs/blender_lego/instant_ngp/2022-07-05_200513/checkpoints\",\n", + " \"trainer.resume_train.load_step=5000\",\n", + " \"viewer.render_image_height=50\",\n", " ],\n", " )\n", " config = setup_config(config)" @@ -110,10 +86,18 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 4, "id": "5d231e36-dd38-489b-a3a3-860cdf49e551", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "2022-07-05 23:33:06.262342: I tensorflow/core/util/util.cc:169] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n" + ] + } + ], "source": [ "trainer = Trainer(config, local_rank=0, world_size=1)\n", "trainer.setup()\n", @@ -122,69 +106,7 @@ }, { "cell_type": "code", - "execution_count": null, - "id": "1aba9ecb-0022-4896-833d-577abe532d23", - "metadata": {}, - "outputs": [], - "source": [ - "aabb = trainer.dataset_inputs_train.scene_bounds.aabb\n", - "vis_utils.draw_aabb(trainer.vis, aabb, name=\"dataset_inputs_train/scene_bounds/aabb\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "4b23fc89-bfc3-40c4-abba-fad1317aeb21", - "metadata": {}, - "outputs": [], - "source": [] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "5a829d31-4fc6-4908-963d-7326d3aa52da", - "metadata": {}, - "outputs": [], - "source": [ - "outputs = trainer.render_image_in_viewer()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "84bd64de-2e0c-458e-a4b7-1084f1cca1c5", - "metadata": {}, - "outputs": [], - "source": [ - "trainer.vis.delete()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "6ed9cdfa-6359-4a70-95c7-033e853e9b50", - "metadata": {}, - "outputs": [], - "source": [ - "while True:\n", - " outputs = trainer.render_image_in_viewer()\n", - " image = outputs[\"rgb_fine\"].cpu().numpy().astype(\"float64\") * 255\n", - " trainer.vis[\"/Cameras/Main Camera\"].set_image(image)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "b3bcd221-05a8-494e-b817-3e2de7f5a0df", - "metadata": {}, - "outputs": [], - "source": [ - "image.shape" - ] - }, - { - "cell_type": "code", - "execution_count": null, + "execution_count": 5, "id": "642b0a25-0e88-4eac-b661-a57a74fb10fc", "metadata": {}, "outputs": [], @@ -194,7 +116,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 6, "id": "938f090b-3b28-4fd6-82dc-d4b5c5917134", "metadata": {}, "outputs": [], @@ -205,10 +127,19 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 13, "id": "fd04703a-aa50-45e4-95fd-4d881ff18d61", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "estimated_time: 3.3333333333333335\n", + "time_elapsed: 14.046947002410889\n" + ] + } + ], "source": [ "# move a camera around\n", "num_cameras = len(dataset_inputs.intrinsics)\n", @@ -228,27 +159,39 @@ "\n", "start = time.time()\n", "for camera in cameras:\n", - " vis_utils.set_camera(vis, camera)\n", + " vis_utils.set_camera(trainer.vis, camera)\n", + "\n", + " # render an image\n", + " outputs = trainer.render_image_in_viewer()\n", + "\n", " time.sleep(1 / fps)\n", + " # break\n", "time_elapsed = time.time() - start\n", "print(\"time_elapsed:\", time_elapsed)" ] }, + { + "cell_type": "code", + "execution_count": 15, + "id": "fbb89f42-bf37-4786-a865-64c4ef018005", + "metadata": {}, + "outputs": [], + "source": [ + "outputs = trainer.render_image_in_viewer()" + ] + }, { "cell_type": "code", "execution_count": null, - "id": "f1cd7183-c514-4f87-a8cc-e14228228e24", + "id": "d91dfa73-ba0c-4412-a908-74ea51c37c58", "metadata": {}, "outputs": [], "source": [] } ], "metadata": { - "interpreter": { - "hash": "4a749acc7d255f078aee52e0584cc77b3cb5aaed1b3a7407ec4262c1bf6cb526" - }, "kernelspec": { - "display_name": "Python 3 (ipykernel)", + "display_name": "Python 3.8.13 ('mattport')", "language": "python", "name": "python3" }, @@ -263,6 +206,11 @@ "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.8.13" + }, + "vscode": { + "interpreter": { + "hash": "f0938bb18d5c7de3ca7c6a73cb6be5f9b7603b44d115fc3e0726e01396d4d2be" + } } }, "nbformat": 4, diff --git a/pyrad/engine/trainer.py b/pyrad/engine/trainer.py index 2be34c032..76498ea9b 100644 --- a/pyrad/engine/trainer.py +++ b/pyrad/engine/trainer.py @@ -17,16 +17,13 @@ Code to train model. """ import logging import os -import random from typing import Dict import torch import torch.distributed as dist -import umsgpack from torch.nn.parallel import DistributedDataParallel as DDP from torchtyping import TensorType -from pyrad.cameras.cameras import get_camera, get_intrinsics_from_intrinsics_matrix from pyrad.cameras.rays import RayBundle from pyrad.data.dataloader import EvalDataloader, setup_dataset_eval, setup_dataset_train from pyrad.graphs.base import setup_graph @@ -36,8 +33,6 @@ from pyrad.utils.config import Config from pyrad.utils.decorators import check_main_thread from pyrad.utils.writer import EventName, TimeWriter from pyrad.viewer.backend import vis_utils -from pyrad.viewer.backend.utils import get_intrinsics_matrix_and_camera_to_world_h -from pyrad.viewer.backend.visualizer import ViewerWindow, Visualizer logging.getLogger("PIL").setLevel(logging.WARNING) @@ -67,15 +62,8 @@ class Trainer: # logging variables writer.setup_event_writers(config.logging, max_iter=config.trainer.max_num_iterations) profiler.setup_profiler(config.logging) - - self.vis = None - if self.config.viewer.enable: - window = ViewerWindow(zmq_url=self.config.viewer.zmq_url) - self.vis = Visualizer(window=window) - logging.info("Connected to viewer at %s", self.config.viewer.zmq_url) - self.vis.delete() - else: - logging.info("Continuing without viewer.") + # visualizer variable + self.visualizer_state = vis_utils.VisualizerState(config.viewer) def setup(self, test_mode=False): """Setup the Trainer by calling other setup functions.""" @@ -106,10 +94,9 @@ class Trainer: def train(self) -> None: """Train the model.""" - - if self.vis: - self.draw_scene_in_viewer() - + self.visualizer_state.init_scene( + image_dataset=self.dataloader_train.image_sampler.dataset, dataset_inputs=self.dataset_inputs_train + ) with TimeWriter(writer, EventName.TOTAL_TRAIN_TIME): num_iterations = self.config.trainer.max_num_iterations iter_dataloader_train = iter(self.dataloader_train) @@ -125,16 +112,10 @@ class Trainer: writer.put_dict(name="Loss/train-loss_dict", scalar_dict=loss_dict, step=step) if step != 0 and self.config.trainer.steps_per_save and step % self.config.trainer.steps_per_save == 0: self._save_checkpoint(self.config.trainer.model_dir, step) - if ( - self.vis - and step != 0 - and self.config.viewer.steps_per_render_image - and step % self.config.viewer.steps_per_render_image == 0 - ): - _ = self.render_image_in_viewer() if step % self.config.trainer.steps_per_test == 0: self.eval_with_dataloader(self.dataloader_eval, step=step) self._write_out_storage(step) + self.visualizer_state.update_scene(step, self.graph) self._write_out_storage(num_iterations) @@ -226,60 +207,6 @@ class Trainer: self.graph.train() return psnr - def draw_scene_in_viewer(self): - """Draw some images and the scene aabb in the viewer.""" - image_dataset_train = self.dataloader_train.image_sampler.dataset - dataset_inputs = self.dataset_inputs_train - indices = random.sample(range(len(image_dataset_train)), k=10) - for idx in indices: - image = image_dataset_train[idx]["image"] - camera = get_camera(dataset_inputs.intrinsics[idx], dataset_inputs.camera_to_world[idx], None) - pose = camera.get_camera_to_world().double().numpy() - K = camera.get_intrinsics_matrix().double().numpy() - vis_utils.draw_camera_frustum( - self.vis, - image=(image.double().numpy() * 255.0), - pose=pose, - K=K, - height=1.0, - name=f"image_dataset_train/{idx:06d}", - displayed_focal_length=0.5, - realistic=False, - ) - aabb = dataset_inputs.scene_bounds.aabb - vis_utils.draw_aabb(self.vis, aabb, name="dataset_inputs_train/scene_bounds/aabb") - - @profiler.time_function - def render_image_in_viewer(self): - """ - Draw an image using the current camera pose from the viewer. - The image is sent of a TCP connection and then uses WebRTC to send it to the viewer. - """ - data = self.vis["/Cameras/Main Camera"].get_object() - message = umsgpack.unpackb(data) - camera_object = message["object"]["object"] - image_height = self.config.viewer.render_image_height - intrinsics_matrix, camera_to_world_h = get_intrinsics_matrix_and_camera_to_world_h( - camera_object, image_height=image_height - ) - - camera_to_world = camera_to_world_h[:3, :] - intrinsics = get_intrinsics_from_intrinsics_matrix(intrinsics_matrix) - camera = get_camera(intrinsics, camera_to_world) - camera_ray_bundle = camera.get_camera_ray_bundle(device=self.device) - camera_ray_bundle.num_rays_per_chunk = self.config.viewer.num_rays_per_chunk - - self.graph.eval() - outputs = self.graph.get_outputs_for_camera_ray_bundle(camera_ray_bundle) - self.graph.train() - - # gross hack to get the image key, depending on which keys the graph uses - rgb_key = "rgb" if "rgb" in outputs else "rgb_fine" - # TODO: make it such that the TCP connection doesn't need float64 - image = outputs[rgb_key].cpu().numpy().astype("float64") * 255 - self.vis["/Cameras/Main Camera"].set_image(image) - return outputs - def eval_with_dataloader(self, dataloader: EvalDataloader, step: int = None) -> None: """Run evaluation with a given dataloader. diff --git a/pyrad/graphs/base.py b/pyrad/graphs/base.py index ed133d0df..4ace7beaf 100644 --- a/pyrad/graphs/base.py +++ b/pyrad/graphs/base.py @@ -17,6 +17,7 @@ The Graph module contains all trainable parameters. """ from abc import abstractmethod from collections import defaultdict +import time from typing import Any, Dict, List, Union import torch @@ -32,6 +33,7 @@ from pyrad.graphs.modules.ray_generator import RayGenerator from pyrad.utils import profiler from pyrad.utils.config import GraphConfig from pyrad.utils.misc import get_masked_dict, instantiate_from_dict_config, is_not_none +from pyrad.viewer.backend.vis_utils import CameraChangeException @profiler.time_function @@ -101,6 +103,8 @@ class Graph(AbstractGraph): self.populate_fields() self.populate_misc_modules() # populate the modules self.callbacks = None + # variable for visualizer to fetch TODO(figure out if there is cleaner way to do this) + self.vis_outputs = None def register_callbacks(self): # pylint:disable=no-self-use """Option to register callback for training functions""" @@ -173,21 +177,27 @@ class Graph(AbstractGraph): @torch.no_grad() def get_outputs_for_camera_ray_bundle(self, camera_ray_bundle: RayBundle): """Takes in camera parameters and computes the output of the graph.""" - assert is_not_none(camera_ray_bundle.num_rays_per_chunk) - image_height, image_width = camera_ray_bundle.origins.shape[:2] - num_rays = len(camera_ray_bundle) - outputs = {} - outputs_lists = defaultdict(list) - for i in range(0, num_rays, camera_ray_bundle.num_rays_per_chunk): - start_idx = i - end_idx = i + camera_ray_bundle.num_rays_per_chunk - ray_bundle = camera_ray_bundle.get_row_major_sliced_ray_bundle(start_idx, end_idx) - outputs = self.forward_after_ray_generator(ray_bundle) - for output_name, output in outputs.items(): - outputs_lists[output_name].append(output) - for output_name, outputs_list in outputs_lists.items(): - outputs[output_name] = torch.cat(outputs_list).view(image_height, image_width, -1) - return outputs + try: + assert is_not_none(camera_ray_bundle.num_rays_per_chunk) + image_height, image_width = camera_ray_bundle.origins.shape[:2] + num_rays = len(camera_ray_bundle) + outputs = {} + outputs_lists = defaultdict(list) + for i in range(0, num_rays, camera_ray_bundle.num_rays_per_chunk): + start_idx = i + end_idx = i + camera_ray_bundle.num_rays_per_chunk + ray_bundle = camera_ray_bundle.get_row_major_sliced_ray_bundle(start_idx, end_idx) + outputs = self.forward_after_ray_generator(ray_bundle) + for output_name, output in outputs.items(): + outputs_lists[output_name].append(output) + time.sleep(0.001) # visualizer allow thread to switch off + for output_name, outputs_list in outputs_lists.items(): + outputs[output_name] = torch.cat(outputs_list).view(image_height, image_width, -1) + time.sleep(0.001) # visualizer allow thread to switch off + self.vis_outputs = outputs + return outputs + except CameraChangeException: + return None def get_outputs_for_camera(self, camera: Camera): """Get the graph outputs for a Camera.""" diff --git a/pyrad/utils/config.py b/pyrad/utils/config.py index e5ca0622b..8b4fa1afb 100644 --- a/pyrad/utils/config.py +++ b/pyrad/utils/config.py @@ -104,6 +104,17 @@ class GraphConfig: field_implementation: Optional[str] = "torch" +@dataclass +class ViewerConfig: + """Configuration for viewer instantiation""" + + enable: bool = MISSING + zmq_url: str = MISSING + min_render_image_height: int = MISSING + max_render_image_height: int = MISSING + num_rays_per_chunk: int = MISSING + + @dataclass class Config: """Full config contents""" @@ -138,7 +149,7 @@ def setup_config(config: DictConfig) -> Config: data = DataConfig(**config.data) graph = GraphConfig(**config.graph) optimizers = config.optimizers - viewer = config.viewer + viewer = ViewerConfig(**config.viewer) return Config( machine=machine, logging=logging, diff --git a/pyrad/viewer/app/package.json b/pyrad/viewer/app/package.json index 56554c16b..7eba6c8e6 100644 --- a/pyrad/viewer/app/package.json +++ b/pyrad/viewer/app/package.json @@ -13,7 +13,7 @@ "react-dom": "^18.1.0", "react-scripts": "^5.0.1", "socket.io-client": "^4.5.1", - "three": "^0.141.0", + "three": "^0.142.0", "websocket": "^1.0.34" }, "scripts": { diff --git a/pyrad/viewer/app/src/SceneNode.js b/pyrad/viewer/app/src/SceneNode.js index cd0cb8e7f..293def860 100644 --- a/pyrad/viewer/app/src/SceneNode.js +++ b/pyrad/viewer/app/src/SceneNode.js @@ -1,137 +1,144 @@ -import * as THREE from 'three'; +import * as THREE from "three"; -const dat = require('dat.gui'); +const dat = require("dat.gui"); function remove_folders(gui) { - for (let name of Object.keys(gui.__folders)) { - let folder = gui.__folders[name]; - remove_folders(folder); - dat.dom.dom.unbind(window, 'resize', folder.__resizeHandler); - gui.removeFolder(folder); - } + for (let name of Object.keys(gui.__folders)) { + let folder = gui.__folders[name]; + remove_folders(folder); + dat.dom.dom.unbind(window, "resize", folder.__resizeHandler); + gui.removeFolder(folder); + } } function dispose(object) { - if (!object) { - return; - } - if (object.geometry) { - object.geometry.dispose(); - } - if (object.material) { - if (Array.isArray(object.material)) { - for (let material of object.material) { - if (material.map) { - material.map.dispose(); - } - material.dispose(); - } - } else { - if (object.material.map) { - object.material.map.dispose(); - } - object.material.dispose(); + if (!object) { + return; + } + if (object.geometry) { + object.geometry.dispose(); + } + if (object.material) { + if (Array.isArray(object.material)) { + for (let material of object.material) { + if (material.map) { + material.map.dispose(); } + material.dispose(); + } + } else { + if (object.material.map) { + object.material.map.dispose(); + } + object.material.dispose(); } + } } export class SceneNode { - constructor(object, folder) { - this.object = object; - this.folder = folder; - this.children = {}; - this.create_controls(); - for (let c of this.object.children) { - this.add_child(c); - } + constructor(object, folder) { + this.object = object; + this.folder = folder; + this.children = {}; + this.create_controls(); + for (let c of this.object.children) { + this.add_child(c); } + } - add_child(object) { - let f = this.folder.addFolder(object.name); - let node = new SceneNode(object, f); - this.children[object.name] = node; - return node; - } + add_child(object) { + let f = this.folder.addFolder(object.name); + let node = new SceneNode(object, f); + this.children[object.name] = node; + return node; + } - create_child(name) { - let obj = new THREE.Group(); - obj.name = name; - this.object.add(obj); - return this.add_child(obj); - } + create_child(name) { + let obj = new THREE.Group(); + obj.name = name; + this.object.add(obj); + return this.add_child(obj); + } - find(path) { - if (path.length == 0) { - return this; - } else { - let name = path[0]; - let child = this.children[name]; - if (child === undefined) { - child = this.create_child(name); - } - return child.find(path.slice(1)); - } + find(path) { + if (path.length == 0) { + return this; + } else { + let name = path[0]; + let child = this.children[name]; + if (child === undefined) { + child = this.create_child(name); + } + return child.find(path.slice(1)); } + } - create_controls() { - if (this.vis_controller !== undefined) { - this.folder.domElement.removeChild(this.vis_controller.domElement); - } - this.vis_controller = new dat.controllers.BooleanController(this.object, "visible"); - this.folder.domElement.prepend(this.vis_controller.domElement); - this.vis_controller.domElement.style.height = "0"; - this.vis_controller.domElement.style.float = "right"; + create_controls() { + if (this.vis_controller !== undefined) { + this.folder.domElement.removeChild(this.vis_controller.domElement); } + this.vis_controller = new dat.controllers.BooleanController( + this.object, + "visible" + ); + this.folder.domElement.prepend(this.vis_controller.domElement); + this.vis_controller.domElement.style.height = "0"; + this.vis_controller.domElement.style.float = "right"; + } - set_property(property, value) { - if (property === "position") { - this.object.position.set(value[0], value[1], value[2]); - } else if (property === "quaternion") { - this.object.quaternion.set(value[0], value[1], value[2], value[3]); - } else if (property === "scale") { - this.object.scale.set(value[0], value[1], value[2]); - } else { - this.object[property] = value; - } - this.vis_controller.updateDisplay(); + set_property(property, value) { + if (property === "position") { + this.object.position.set(value[0], value[1], value[2]); + } else if (property === "quaternion") { + this.object.quaternion.set(value[0], value[1], value[2], value[3]); + } else if (property === "scale") { + this.object.scale.set(value[0], value[1], value[2]); + } else { + this.object[property] = value; } + this.vis_controller.updateDisplay(); + } - set_transform(matrix) { - let mat = new THREE.Matrix4(); - mat.fromArray(matrix); - mat.decompose(this.object.position, this.object.quaternion, this.object.scale); - } + set_transform(matrix) { + let mat = new THREE.Matrix4(); + mat.fromArray(matrix); + mat.decompose( + this.object.position, + this.object.quaternion, + this.object.scale + ); + } - set_object(object) { - let parent = this.object.parent; - this.dispose_recursive(); - this.object.parent.remove(this.object); - this.object = object; - parent.add(object); - this.create_controls(); - } + set_object(object) { + let parent = this.object.parent; + this.dispose_recursive(); + this.object.parent.remove(this.object); + this.object = object; + parent.add(object); + this.create_controls(); + } - dispose_recursive() { - for (let name of Object.keys(this.children)) { - this.children[name].dispose_recursive(); - } - dispose(this.object); + dispose_recursive() { + for (let name of Object.keys(this.children)) { + this.children[name].dispose_recursive(); } + dispose(this.object); + } - delete(path) { - if (path.length == 0) { - console.error("Can't delete an empty path"); - } else { - let parent = this.find(path.slice(0, path.length - 1)); - let name = path[path.length - 1]; - let child = parent.children[name]; - if (child !== undefined) { - child.dispose_recursive(); - parent.object.remove(child.object); - remove_folders(child.folder); - parent.folder.removeFolder(child.folder); - delete parent.children[name]; - } - } + delete(path) { + if (path.length == 0) { + console.error("Can't delete an empty path"); + } else { + let parent = this.find(path.slice(0, path.length - 1)); + let name = path[path.length - 1]; + let child = parent.children[name]; + if (child !== undefined) { + child.dispose_recursive(); + parent.object.remove(child.object); + remove_folders(child.folder); + parent.folder.removeFolder(child.folder); + delete parent.children[name]; + } } -} \ No newline at end of file + } +} diff --git a/pyrad/viewer/app/src/Viewer.js b/pyrad/viewer/app/src/Viewer.js index 6702287c7..905ca0d90 100644 --- a/pyrad/viewer/app/src/Viewer.js +++ b/pyrad/viewer/app/src/Viewer.js @@ -74,8 +74,19 @@ export class Viewer extends Component { this.state.controls_main.update(); } this.state.renderer_main.render(this.state.scene, camera_main); + this.send_camera_over_websocket(); + } + + send_camera_over_websocket() { + /* update the camera information in the python server + if the websocket is connected */ + console.log("send_camera_over_websocket"); if (this.state.websocket.readyState === WebSocket.OPEN) { // update the camera information in the python server + let camera_main = findCameraObjectUnderObject3D( + this.state.scene_tree.find(["Cameras", "Main Camera"]).object + ); + console.log(camera_main.toJSON()); let cmd = "set_object"; let path = "Cameras/Main Camera"; let data = { @@ -121,7 +132,12 @@ export class Viewer extends Component { this.state.controls_main.staticMoving = false; // false is default this.state.controls_main.target.set(0, 0, 0); // focus point of the controls this.state.controls_main.autoRotate = false; + this.state.controls_main.dynamicDampingFactor = 1.0; this.state.controls_main.update(); + // this.send_camera_over_websocket(); + // this.state.controls_main.addEventListener("change", () => { + // this.send_camera_over_websocket(); + // }); } } @@ -190,6 +206,9 @@ export class Viewer extends Component { let answer = cmd.data; this.state.pc.setRemoteDescription(answer); } + + // possibly update the camera information in the python server + // this.send_camera_over_websocket(); } getViewportWidth() { diff --git a/pyrad/viewer/app/yarn.lock b/pyrad/viewer/app/yarn.lock index d1e46ffb9..7af92d056 100644 --- a/pyrad/viewer/app/yarn.lock +++ b/pyrad/viewer/app/yarn.lock @@ -8745,10 +8745,10 @@ text-table@^0.2.0: resolved "https://registry.yarnpkg.com/text-table/-/text-table-0.2.0.tgz#7f5ee823ae805207c00af2df4a84ec3fcfa570b4" integrity sha512-N+8UisAXDGk8PFXP4HAzVR9nbfmVJ3zYLAWiTIoqC5v5isinhr+r5uaO8+7r3BMfuNIufIsA7RdpVgacC2cSpw== -three@^0.141.0: - version "0.141.0" - resolved "https://registry.yarnpkg.com/three/-/three-0.141.0.tgz#16677a12b9dd0c3e1568ebad0fd09de15d5a8216" - integrity sha512-JaSDAPWuk4RTzG5BYRQm8YZbERUxTfTDVouWgHMisS2to4E5fotMS9F2zPFNOIJyEFTTQDDKPpsgZVThKU3pXA== +three@^0.142.0: + version "0.142.0" + resolved "https://registry.yarnpkg.com/three/-/three-0.142.0.tgz#89e226a16221f212eb1d40f0786604b711f28aed" + integrity sha512-ESjPO+3geFr+ZUfVMpMnF/eVU2uJPOh0e2ZpMFqjNca1wApS9lJb7E4MjwGIczgt9iuKd8PEm6Pfgp2bJ92Xtg== throat@^6.0.1: version "6.0.1" diff --git a/pyrad/viewer/backend/run_zmq_server.py b/pyrad/viewer/backend/run_zmq_server.py index f8ad6478a..913e737b0 100644 --- a/pyrad/viewer/backend/run_zmq_server.py +++ b/pyrad/viewer/backend/run_zmq_server.py @@ -210,7 +210,8 @@ class ZMQWebSocketBridge(object): path = list(filter(lambda x: len(x) > 0, frames[1].decode("utf-8").split("/"))) data = frames[2] if cmd in WEBSOCKET_COMMANDS: - self.forward_to_websockets(frames) + if cmd != "get_object": + self.forward_to_websockets(frames) if cmd == "set_transform": find_node(self.tree, path).transform = data elif cmd == "set_object": diff --git a/pyrad/viewer/backend/vis_utils.py b/pyrad/viewer/backend/vis_utils.py index e4770d6d5..ca56e4daf 100644 --- a/pyrad/viewer/backend/vis_utils.py +++ b/pyrad/viewer/backend/vis_utils.py @@ -16,17 +16,187 @@ """ import copy +import logging import random +import sys +import threading +import time -import cv2 import matplotlib.pyplot as plt import numpy as np -from pyrad.cameras.cameras import Camera +import torch +import umsgpack +from omegaconf import DictConfig +from pyrad.cameras.rays import RayBundle +from pyrad.utils.config import ViewerConfig import pyrad.viewer.backend.cameras as c import pyrad.viewer.backend.geometry as g import pyrad.viewer.backend.transformations as tf +from pyrad.cameras.cameras import Camera, get_camera, get_intrinsics_from_intrinsics_matrix +from pyrad.utils import profiler from pyrad.viewer.backend import ViewerWindow, Visualizer +from pyrad.viewer.backend.utils import get_intrinsics_matrix_and_camera_to_world_h + + +class CameraChangeException(Exception): + """Basic camera exception to interrupt visualizer""" + + pass + + +class SetTrace: + """Basic trace function""" + + def __init__(self, func): + self.func = func + + def __enter__(self): + sys.settrace(self.func) + return self + + def __exit__(self, ext_type, exc_value, traceback): + sys.settrace(None) + + +class VisualizerState: + """Class to hold state for""" + + def __init__(self, config: ViewerConfig): + self.config = config + self.vis = None + if self.config.enable: + window = ViewerWindow(zmq_url=self.config.zmq_url) + self.vis = Visualizer(window=window) + logging.info("Connected to viewer at %s", self.config.zmq_url) + self.vis.delete() + else: + logging.info("Continuing without viewer.") + + # visualizer specific variables + self.prev_camera_matrix = None + self.res_upscale_factor = 1 + self.lock = threading.Lock() + self.check_interrupt_vis = False + self.check_done_render = True + + def init_scene(self, image_dataset, dataset_inputs): + """initializes the scene with the datasets""" + if self.vis: + self._draw_scene_in_viewer(image_dataset, dataset_inputs) + + def update_scene(self, step, graph): + """updates the scene based on the graph weights""" + if self._is_render_step(step): + self._render_image_in_viewer(graph) + + def _check_interrupt(self, frame, event, arg): + if event == "line": + if self.check_interrupt_vis: + raise CameraChangeException + return self._check_interrupt + + def _is_render_step(self, step, default_steps=10): + """dynamically calculate when to render grapic based on resolution of image""" + if self.vis and step != 0: + steps_per_render_image = min(default_steps * self.res_upscale_factor, 100) + return step % steps_per_render_image == 0 + return False + + def _draw_scene_in_viewer(self, image_dataset, dataset_inputs): + """Draw some images and the scene aabb in the viewer.""" + indices = random.sample(range(len(image_dataset)), k=10) + for idx in indices: + image = image_dataset[idx]["image"] + camera = get_camera(dataset_inputs.intrinsics[idx], dataset_inputs.camera_to_world[idx], None) + pose = camera.get_camera_to_world().double().numpy() + K = camera.get_intrinsics_matrix().double().numpy() + draw_camera_frustum( + self.vis, + image=(image.double().numpy() * 255.0), + pose=pose, + K=K, + height=1.0, + name=f"image_dataset/{idx:06d}", + displayed_focal_length=0.5, + realistic=False, + ) + aabb = dataset_inputs.scene_bounds.aabb + draw_aabb(self.vis, aabb, name="dataset_inputs_train/scene_bounds/aabb") + + def _async_check_camera_update(self): + """Async function to check whether camera has been updated in visualizer""" + with self.lock: + self.check_done_render = False + while not self.check_done_render: + data = self.vis["/Cameras/Main Camera"].get_object() + message = umsgpack.unpackb(data) + camera_object = message["object"]["object"] + if self.prev_camera_matrix is None or not np.array_equal(camera_object["matrix"], self.prev_camera_matrix): + with self.lock: + self.check_interrupt_vis = True + self.prev_camera_matrix = camera_object["matrix"] + self.res_upscale_factor = 1 + time.sleep(0.001) + + @torch.no_grad() + def _async_get_visualizer_outputs(self, graph, camera_ray_bundle: RayBundle): + """async getter function for visualizer without returning""" + with SetTrace(self._check_interrupt): + graph.get_outputs_for_camera_ray_bundle(camera_ray_bundle) + with self.lock: + self.check_done_render = True + self.check_interrupt_vis = False + + @profiler.time_function + def _render_image_in_viewer(self, graph): + """ + Draw an image using the current camera pose from the viewer. + The image is sent of a TCP connection and then uses WebRTC to send it to the viewer. + """ + data = self.vis["/Cameras/Main Camera"].get_object() + message = umsgpack.unpackb(data) + camera_object = message["object"]["object"] + # hacky way to prevent overflow check to see if < 100; TODO(make less hacky) + if self.prev_camera_matrix is not None and np.array_equal(camera_object["matrix"], self.prev_camera_matrix): + self.res_upscale_factor = min(self.res_upscale_factor * 2, 100) + else: + self.prev_camera_matrix = camera_object["matrix"] + self.res_upscale_factor = 1 + + image_height = min( + self.config.min_render_image_height * self.res_upscale_factor, + self.config.max_render_image_height, + ) + intrinsics_matrix, camera_to_world_h = get_intrinsics_matrix_and_camera_to_world_h( + camera_object, image_height=image_height + ) + + camera_to_world = camera_to_world_h[:3, :] + intrinsics = get_intrinsics_from_intrinsics_matrix(intrinsics_matrix) + camera = get_camera(intrinsics, camera_to_world) + camera_ray_bundle = camera.get_camera_ray_bundle(device=graph.get_device()) + camera_ray_bundle.num_rays_per_chunk = self.config.num_rays_per_chunk + + graph.eval() + try: + check_thread = threading.Thread(target=self._async_check_camera_update) + render_thread = threading.Thread(target=self._async_get_visualizer_outputs, args=(graph, camera_ray_bundle)) + check_thread.start() + render_thread.start() + check_thread.join() + render_thread.join() + except Exception as e: # pylint: disable=broad-except + print(e) + graph.train() + outputs = graph.vis_outputs + if outputs is not None: + # gross hack to get the image key, depending on which keys the graph uses + rgb_key = "rgb" if "rgb" in outputs else "rgb_fine" + # TODO: make it such that the TCP connection doesn't need float64 + image = outputs[rgb_key].cpu().numpy().astype("float64") * 255 + self.vis["/Cameras/Main Camera"].set_image(image) + return outputs def get_vis(zmq_url="tcp://0.0.0.0:6000"): @@ -47,50 +217,6 @@ def get_random_color(): return color -def plot_correspondences(pair, plot=True): - """Draw what the pair looks like, with the correspondences as lines.""" - image0 = (pair["data0"]["image"] * 255).astype("uint8").transpose((1, 2, 0)) - image1 = (pair["data1"]["image"] * 255).astype("uint8").transpose((1, 2, 0)) - original_image = np.hstack([image0, image1]) - h, w, _ = image0.shape - - matches_image = original_image.copy() - # draw lines - correspondences = list(pair["correspondences"]) - num = min(20, len(correspondences)) - for x0, y0, x1, y1 in np.array(random.sample(correspondences, k=num)).astype("uint64"): - color = get_random_color() - thickness = 2 - matches_image = cv2.line(matches_image, (x0, y0), (int(w + x1), y1), color, thickness) - # show image - if plot: - print("Correspondences:") - plt.figure(figsize=(20, 10)) - plt.imshow(matches_image) - plt.show() - - # ---------- - # example training data - equal_distances_image = original_image.copy() - for i in range(num): - c0, c1 = np.array(random.sample(correspondences, k=2)).astype("uint64") - color = get_random_color() - thickness = 2 - p0, p1 = c0[:2], c1[:2] - equal_distances_image = cv2.line(equal_distances_image, tuple(p0), tuple(p1), color, thickness) - p0, p1 = c0[2:], c1[2:] - p0[0] += w - p1[0] += w - equal_distances_image = cv2.line(equal_distances_image, tuple(p0), tuple(p1), color, thickness) - if plot: - print("Example training data:") - plt.figure(figsize=(20, 10)) - plt.imshow(equal_distances_image) - plt.show() - - return matches_image, equal_distances_image - - def show_ply(vis, ply_path, name="ply", color=None): """Show the PLY file in the 3D viewer. Specify the full filename as input.""" assert ply_path.endswith(".ply") @@ -171,17 +297,6 @@ def draw_camera_frustum( vis[name].set_transform(pose) -def set_camera_render(vis, intrinsics=None, pose=None, name="renderer"): - """Place a three.js camera in the scene. - This can be used to render an image from. - """ - full_name_str = f"/Cameras/{name}/rotated" - g_camera = c.PerspectiveCamera(fov=120, aspect=1.0, near=0.01, far=1000) - g_camera_helper = c.CameraHelper(g_camera) - # vis[full_name_str].set_object(g_camera) - vis[full_name_str].set_object(g_camera_helper) - - def set_persp_camera(vis, pose, K, colmap=True): """Assumes simple pinhole model for intrinsics. Args: @@ -196,21 +311,9 @@ def set_persp_camera(vis, pose, K, colmap=True): focal_length = K[0, 0] x = pp_h / (focal_length) fov = 2.0 * np.arctan(x) * (180.0 / np.pi) - vis["/Cameras/Main Camera R/"].set_property("fov", fov) - vis["/Cameras/Main Camera R/"].set_property("aspect", float(pp_w / pp_h)) # three.js expects width/height - vis["/Cameras/Main Camera R"].set_transform(pose_processed) - - -def set_orth_camera(vis, pose, width, height, colmap=True): - """ """ - pose_processed = copy.deepcopy(pose) - if colmap: - pose_processed[:, 1:3] *= -1 - vis["/Cameras/Main Camera Orth"].set_transform(pose_processed) - vis["/Cameras/Main Camera Orth/"].set_property("left", -width / 2.0) - vis["/Cameras/Main Camera Orth/"].set_property("right", width / 2.0) - vis["/Cameras/Main Camera Orth/"].set_property("top", height / 2.0) - vis["/Cameras/Main Camera Orth/"].set_property("bottom", -height / 2.0) + vis["/Cameras/Main Camera/"].set_property("fov", fov) + vis["/Cameras/Main Camera/"].set_property("aspect", float(pp_w / pp_h)) # three.js expects width/height + vis["/Cameras/Main Camera/"].set_transform(pose_processed) def set_camera(vis, camera: Camera):