From 25a1f91deed0560d3177028554ec30e4bfc75eca Mon Sep 17 00:00:00 2001 From: evonneng Date: Tue, 12 Jul 2022 18:52:04 -0700 Subject: [PATCH] Visualizer: toggle between different output options (#122) * added dynamic setting options based on graph outputs * added functionality to view different types of outputs * fixed first value set to nothing * moved try except in smaller scope --- pyrad/viewer/app/src/Viewer.jsx | 30 ++++++++++++++++ pyrad/viewer/server/commands.py | 22 ++++++++++++ pyrad/viewer/server/server.py | 18 +++++++++- pyrad/viewer/server/viewer_utils.py | 56 +++++++++++++++++++++++------ pyrad/viewer/server/visualizer.py | 21 ++++++++++- 5 files changed, 135 insertions(+), 12 deletions(-) diff --git a/pyrad/viewer/app/src/Viewer.jsx b/pyrad/viewer/app/src/Viewer.jsx index 68d404913..f1ef01220 100644 --- a/pyrad/viewer/app/src/Viewer.jsx +++ b/pyrad/viewer/app/src/Viewer.jsx @@ -97,6 +97,22 @@ export class Viewer extends Component { } } + send_output_type_over_websocket(value) { + /* update the output option in the python server + if the user changes selection */ + if (this.state.websocket.readyState === WebSocket.OPEN) { + let cmd = "set_output_type"; + let path = "Output Type"; + let data = { + type: cmd, + path: path, + output_type: value, + }; + let message = msgpack.encode(data); + this.state.websocket.send(message); + } + } + set_object(path, object) { this.state.scene_tree.find(path.concat([""])).set_object(object); // add controls if object is a camera @@ -161,6 +177,18 @@ export class Viewer extends Component { } } + set_output_options(object) { + let output_options_control = new function () { + this.output_options = 'default'; + } + // add controls + this.state.gui.add(output_options_control, 'output_options', object).listen().onChange( + (value) => { + this.send_output_type_over_websocket(value); + } + ) + } + handle_command(cmd) { // convert binary serialization format back to JSON cmd = msgpack.decode(new Uint8Array(cmd)); @@ -179,6 +207,8 @@ export class Viewer extends Component { } else if (cmd.type === "set_property") { let path = split_path(cmd.path); this.set_property(path, cmd.property, cmd.value); + } else if(cmd.type == 'set_output_options') { + this.set_output_options(cmd.output_options); } // web rtc commands else if (cmd.type === "answer") { diff --git a/pyrad/viewer/server/commands.py b/pyrad/viewer/server/commands.py index bfceb3b3d..42631e2ac 100644 --- a/pyrad/viewer/server/commands.py +++ b/pyrad/viewer/server/commands.py @@ -72,6 +72,28 @@ class SetTransform: return {"type": "set_transform", "path": self.path.lower(), "matrix": list(self.matrix.T.flatten())} +class SetOutputOptions: + __slots__ = ["output_options", "path"] + + def __init__(self, output_options, path): + self.output_options = output_options + self.path = path + + def lower(self): + return {"type": "set_output_options", "path": self.path.lower(), "output_options": list(self.output_options)} + + +class SetOutputType: + __slots__ = ["output_type", "path"] + + def __init__(self, output_type, path): + self.output_type = output_type + self.path = path + + def lower(self): + return {"type": "set_output_options", "path": self.path.lower(), "output_options": list(self.output_type)} + + class Delete: __slots__ = ["path"] diff --git a/pyrad/viewer/server/server.py b/pyrad/viewer/server/server.py index 13db0a513..a9d38ff7f 100644 --- a/pyrad/viewer/server/server.py +++ b/pyrad/viewer/server/server.py @@ -37,7 +37,15 @@ MAX_ATTEMPTS = 1000 DEFAULT_ZMQ_METHOD = "tcp" DEFAULT_ZMQ_PORT = 6000 DEFAULT_WEBSOCKET_PORT = 8051 -WEBSOCKET_COMMANDS = ["set_transform", "set_object", "get_object", "set_property", "delete"] +WEBSOCKET_COMMANDS = [ + "set_transform", + "set_object", + "set_output_options", + "set_output_type", + "get_object", + "set_property", + "delete", +] WEBRTC_COMMANDS = ["set_image"] @@ -95,6 +103,10 @@ class WebSocketHandler(tornado.websocket.WebSocketHandler): elif type_ == "set_object": find_node(self.bridge.tree, path).object = data find_node(self.bridge.tree, path).properties = [] + elif type_ == "set_output_options": + find_node(self.bridge.tree, path).object = data + elif type_ == "set_output_type": + find_node(self.bridge.tree, path).object = data elif type_ == "offer": offer = RTCSessionDescription(m["data"]["sdp"], m["data"]["type"]) @@ -177,6 +189,10 @@ class ZMQWebSocketBridge(object): elif cmd == "set_object": find_node(self.tree, path).object = data find_node(self.tree, path).properties = [] + elif cmd == "set_output_options": + find_node(self.tree, path).object = data + elif cmd == "set_output_type": + find_node(self.tree, path).object = data elif cmd == "get_object": data = find_node(self.tree, path).object if isinstance(data, type(None)): diff --git a/pyrad/viewer/server/viewer_utils.py b/pyrad/viewer/server/viewer_utils.py index 056857787..885e40ad3 100644 --- a/pyrad/viewer/server/viewer_utils.py +++ b/pyrad/viewer/server/viewer_utils.py @@ -21,6 +21,7 @@ import random import sys import threading import time +from typing import List import numpy as np import torch @@ -39,7 +40,7 @@ from pyrad.viewer.server.utils import get_intrinsics_matrix_and_camera_to_world_ from pyrad.viewer.server.transformations import get_translation_matrix -class CameraChangeException(Exception): +class IOChangeException(Exception): """Basic camera exception to interrupt visualizer""" @@ -73,10 +74,13 @@ class VisualizerState: # visualizer specific variables self.prev_camera_matrix = None + self.prev_output_type = "rgb" self.res_upscale_factor = 1 self.check_interrupt_vis = False self.check_done_render = True + self.outputs_set = False + def init_scene(self, image_dataset: ImageDataset, dataset_inputs: DatasetInputs) -> None: """initializes the scene with the datasets""" if self.vis: @@ -92,7 +96,7 @@ class VisualizerState: if event == "line": if self.check_interrupt_vis and self.res_upscale_factor > 1: self.res_upscale_factor = 1 - raise CameraChangeException + raise IOChangeException return self._check_interrupt def _is_render_step(self, step: int, default_steps: int = 5) -> bool: @@ -129,16 +133,27 @@ class VisualizerState: K = camera.get_intrinsics_matrix() set_persp_intrinsics_matrix(self.vis, K.double().numpy()) - def _async_check_camera_update(self) -> None: + def _async_check_io_update(self) -> None: """Async function to check whether camera has been updated in visualizer""" self.check_done_render = False while not self.check_done_render: + # check camera data = self.vis["/Cameras/Main Camera"].get_object() if data is None: return camera_object = data["object"]["object"] if self.prev_camera_matrix is None or not np.array_equal(camera_object["matrix"], self.prev_camera_matrix): self.check_interrupt_vis = True + + # check output type + data = self.vis["/Output Type"].get_object() + if data is None: + output_type = "rgb" + else: + output_type = data["output_type"] + if self.prev_output_type != output_type: + self.check_interrupt_vis = True + time.sleep(0.001) @torch.no_grad() @@ -155,6 +170,7 @@ class VisualizerState: 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. """ + # check and perform camera updates data = self.vis["/Cameras/Main Camera"].get_object() if data is None: return @@ -166,6 +182,14 @@ class VisualizerState: self.prev_camera_matrix = camera_object["matrix"] self.res_upscale_factor = 1 + # check and perform output type updates + data = self.vis["/Output Type"].get_object() + if data is None: + output_type = "rgb" + else: + output_type = data["output_type"] + self.prev_output_type = output_type + image_height = min( self.config.min_render_image_height * self.res_upscale_factor, self.config.max_render_image_height, @@ -181,11 +205,11 @@ class VisualizerState: camera_ray_bundle.num_rays_per_chunk = self.config.num_rays_per_chunk graph.eval() + check_thread = threading.Thread(target=self._async_check_io_update) + render_thread = threading.Thread(target=self._async_get_visualizer_outputs, args=(graph, camera_ray_bundle)) + check_thread.start() + render_thread.start() 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 @@ -193,9 +217,16 @@ class VisualizerState: graph.train() outputs = graph.vis_outputs if outputs is not None: + if not self.outputs_set: + set_output_options(self.vis, list(outputs.keys())) + self.outputs_set = True # gross hack to get the image key, depending on which keys the graph uses - rgb_key = "rgb" if "rgb" in outputs else "rgb_fine" - image = (outputs[rgb_key].cpu().numpy() * 255).astype("uint8") + # rgb_key = "rgb" if "rgb" in outputs else "rgb_fine" + # image = (outputs[rgb_key].cpu().numpy() * 255).astype("uint8") + image_output = outputs[output_type].cpu().numpy() * 255 + if image_output.shape[-1] == 1: + image_output = np.tile(image_output, (1, 1, 3)) + image = (image_output).astype("uint8") self.vis["/Cameras/Main Camera"].set_image(image) @@ -204,7 +235,7 @@ def interruptable_get_outputs_for_camera_ray_bundle(graph: Graph, camera_ray_bun try: outputs = graph.get_outputs_for_camera_ray_bundle(camera_ray_bundle) return outputs - except CameraChangeException: + except IOChangeException: return None finally: # TODO(): find out better way to do interrupts for camera change @@ -218,6 +249,11 @@ def get_default_vis() -> Viewer: return viewer +def set_output_options(vis: Viewer, output_options: List[str]): + """Sets the possible list of output options for user to toggle""" + vis["output_options"].set_output_options(output_options) + + def show_box_test(vis: Viewer): """Simple test to draw a box and make sure everything is working.""" vis["box"].set_object(g.Box([1.0, 1.0, 1.0]), material=g.MeshPhongMaterial(color=0xFF0000)) diff --git a/pyrad/viewer/server/visualizer.py b/pyrad/viewer/server/visualizer.py index 98358ffbf..882223a88 100644 --- a/pyrad/viewer/server/visualizer.py +++ b/pyrad/viewer/server/visualizer.py @@ -22,7 +22,15 @@ import numpy as np import umsgpack import zmq -from pyrad.viewer.server.commands import Delete, GetObject, SetObject, SetProperty, SetTransform +from pyrad.viewer.server.commands import ( + Delete, + GetObject, + SetObject, + SetOutputOptions, + SetOutputType, + SetProperty, + SetTransform, +) from pyrad.viewer.server.path import Path @@ -144,6 +152,17 @@ class Viewer(object): def set_property(self, key, value): return self.window.send(SetProperty(key, value, self.path)) + def set_output_options(self, options): + return self.window.send(SetOutputOptions(options, self.path)) + + def set_output_type(self, type): + data = self.window.send(SetOutputType(type, self.path)) + data = umsgpack.unpackb(data) + if isinstance(data, str) and data.find("error") == 0: + # some error meaning that the object does not exist + return None + return data + def delete(self): return self.window.send(Delete(self.path))