From b349467d6195ade91b7c675bd27da605f7dd14e4 Mon Sep 17 00:00:00 2001
From: Brent Yi <yibrenth@gmail.com>
Date: Thu, 12 Oct 2023 17:28:13 +0800
Subject: [PATCH] Refactor for `get_render()`

---
 examples/19_get_renders.py                  |  2 +-
 src/viser/_viser.py                         | 72 ++++++++++-----------
 src/viser/client/src/WebsocketInterface.tsx |  9 +--
 3 files changed, 38 insertions(+), 45 deletions(-)

diff --git a/examples/19_get_renders.py b/examples/19_get_renders.py
index 61275f62e..3f94426c7 100644
--- a/examples/19_get_renders.py
+++ b/examples/19_get_renders.py
@@ -33,7 +33,7 @@ def _(event: viser.GuiEvent) -> None:
                 line_width=3.0,
                 color=onp.random.uniform(size=3),
             )
-            images.append(client.get_render(height=720, width=1280))
+            images.append(client.camera.get_render(height=720, width=1280))
 
         print("Writing GIF...")
         iio.imwrite("saved.gif", images)
diff --git a/src/viser/_viser.py b/src/viser/_viser.py
index b2c1ad936..600c097db 100644
--- a/src/viser/_viser.py
+++ b/src/viser/_viser.py
@@ -169,37 +169,6 @@ def on_update(
         self._state.camera_cb.append(callback)
         return callback
 
-
-@dataclasses.dataclass
-class _ClientHandleState:
-    server: infra.Server
-    connection: infra.ClientConnection
-
-
-@dataclasses.dataclass
-class ClientHandle(MessageApi, GuiApi):
-    """Handle for interacting with a specific client. Can be used to send messages to
-    individual clients and read/write camera information."""
-
-    client_id: int
-    """Unique ID for this client."""
-    camera: CameraHandle
-    """Handle for reading from and manipulating the client's viewport camera."""
-    _state: _ClientHandleState
-
-    def __post_init__(self):
-        super().__init__(self._state.connection)
-
-    @override
-    def _get_api(self) -> MessageApi:
-        """Message API to use."""
-        return self
-
-    @override
-    def _queue_unsafe(self, message: _messages.Message) -> None:
-        """Define how the message API should send messages."""
-        self._state.connection.send(message)
-
     def get_render(
         self, height: int, width: int, transport_format: Literal["png", "jpeg"] = "jpeg"
     ) -> onp.ndarray:
@@ -219,11 +188,13 @@ def get_render(
         render_ready_event = threading.Event()
         out: Optional[onp.ndarray] = None
 
+        connection = self.client._state.connection
+
         def got_render_cb(
             client_id: int, message: _messages.GetRenderResponseMessage
         ) -> None:
             del client_id
-            self._state.connection.unregister_handler(
+            connection.unregister_handler(
                 _messages.GetRenderResponseMessage, got_render_cb
             )
             nonlocal out
@@ -233,10 +204,8 @@ def got_render_cb(
             )
             render_ready_event.set()
 
-        self._state.connection.register_handler(
-            _messages.GetRenderResponseMessage, got_render_cb
-        )
-        self._queue(
+        connection.register_handler(_messages.GetRenderResponseMessage, got_render_cb)
+        self.client._queue(
             _messages.GetRenderRequestMessage(
                 "image/jpeg" if transport_format == "jpeg" else "image/png",
                 height=height,
@@ -251,6 +220,37 @@ def got_render_cb(
         assert out is not None
         return out
 
+
+@dataclasses.dataclass
+class _ClientHandleState:
+    server: infra.Server
+    connection: infra.ClientConnection
+
+
+@dataclasses.dataclass
+class ClientHandle(MessageApi, GuiApi):
+    """Handle for interacting with a specific client. Can be used to send messages to
+    individual clients and read/write camera information."""
+
+    client_id: int
+    """Unique ID for this client."""
+    camera: CameraHandle
+    """Handle for reading from and manipulating the client's viewport camera."""
+    _state: _ClientHandleState
+
+    def __post_init__(self):
+        super().__init__(self._state.connection)
+
+    @override
+    def _get_api(self) -> MessageApi:
+        """Message API to use."""
+        return self
+
+    @override
+    def _queue_unsafe(self, message: _messages.Message) -> None:
+        """Define how the message API should send messages."""
+        self._state.connection.send(message)
+
     @contextlib.contextmanager
     def atomic(self) -> Generator[None, None, None]:
         """Returns a context where:
diff --git a/src/viser/client/src/WebsocketInterface.tsx b/src/viser/client/src/WebsocketInterface.tsx
index 6bd312691..aaf839b4c 100644
--- a/src/viser/client/src/WebsocketInterface.tsx
+++ b/src/viser/client/src/WebsocketInterface.tsx
@@ -695,9 +695,6 @@ export function FrameSynchronizedMessageHandler() {
   const viewer = useContext(ViewerContext)!;
   const messageQueueRef = viewer.messageQueueRef;
 
-  // We'll reuse the same canvas.
-  const renderBufferCanvas = React.useMemo(() => new OffscreenCanvas(1, 1), []);
-
   useFrame(() => {
     // Send a render along if it was requested!
     if (viewer.getRenderRequestState.current === "triggered") {
@@ -709,11 +706,7 @@ export function FrameSynchronizedMessageHandler() {
       const targetHeight = viewer.getRenderRequest.current!.height;
 
       // We'll save a render to an intermediate canvas with the requested dimensions.
-      if (renderBufferCanvas.width !== targetWidth)
-        renderBufferCanvas.width = targetWidth;
-      if (renderBufferCanvas.height !== targetHeight)
-        renderBufferCanvas.height = targetHeight;
-
+      const renderBufferCanvas = new OffscreenCanvas(targetWidth, targetHeight);
       const ctx = renderBufferCanvas.getContext("2d")!;
       ctx.reset();
       // Use a white background for JPEGs, which don't have an alpha channel.