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

@@ -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(["<object>"])).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") {

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"]
 

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)):

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))

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))