From 633f7dd58bd4bf96b172f1b9135485f2f2c35e9e Mon Sep 17 00:00:00 2001 From: Brent Yi Date: Wed, 31 Jul 2024 18:24:20 -0700 Subject: [PATCH] Optimize image transport + decoding --- examples/assets/mdx_example.mdx | 6 +- src/viser/_gui_handles.py | 10 +-- src/viser/_messages.py | 8 +-- src/viser/_scene_api.py | 33 ++++------ src/viser/client/src/App.tsx | 6 +- src/viser/client/src/MessageHandler.tsx | 75 +++++++++++++--------- src/viser/client/src/WebsocketMessages.tsx | 8 +-- 7 files changed, 78 insertions(+), 68 deletions(-) diff --git a/examples/assets/mdx_example.mdx b/examples/assets/mdx_example.mdx index f72a4f6b5..fe73a976d 100644 --- a/examples/assets/mdx_example.mdx +++ b/examples/assets/mdx_example.mdx @@ -16,7 +16,7 @@ In inline code blocks, you can show off colors with color chips: `#FED363` Adding images from a remote origin is simple. -![Viser Logo](http://nerfstudio-project.github.io/viser/_static/viser.svg) +![Viser Logo](https://viser.studio/latest/_static/logo.svg) For local images with relative paths, you can either directly use a [data URL](https://developer.mozilla.org/en-US/docs/Web/HTTP/Basics_of_HTTP/Data_URLs) @@ -30,7 +30,7 @@ Tables follow the standard markdown spec: | Application | Description | | ---------------------------------------------------- | -------------------------------------------------- | -| [Nerfstudio](https://nerf.studio) | A collaboration friendly studio for NeRFs | +| [NS](https://nerf.studio) | A collaboration friendly studio for NeRFs | | [Viser](https://nerfstudio-project.github.io/viser/) | An interactive 3D visualization toolbox for Python | Code blocks, while being not nearly as exciting as some of the things presented, @@ -90,5 +90,3 @@ So that's MDX in Viser. It has support for: blocks, inline code - [x] Color chips - [x] JSX enhanced components -- [ ] Prism highlighted code blocks and code block tabs -- [ ] Exposed Mantine in markdown diff --git a/src/viser/_gui_handles.py b/src/viser/_gui_handles.py index 9196ccd4e..ee5eecd64 100644 --- a/src/viser/_gui_handles.py +++ b/src/viser/_gui_handles.py @@ -1,9 +1,9 @@ from __future__ import annotations +import base64 import dataclasses import re import time -import urllib.parse import uuid import warnings from pathlib import Path @@ -16,7 +16,7 @@ from ._icons import svg_from_icon from ._icons_enum import IconName from ._messages import GuiCloseModalMessage, GuiRemoveMessage, GuiUpdateMessage, Message -from ._scene_api import _encode_image_base64 +from ._scene_api import _encode_image_binary from .infra import ClientId if TYPE_CHECKING: @@ -537,9 +537,9 @@ def _get_data_url(url: str, image_root: Path | None) -> str: image_root = Path(__file__).parent try: image = iio.imread(image_root / url) - data_uri = _encode_image_base64(image, "png") - url = urllib.parse.quote(f"{data_uri[1]}") - return f"data:{data_uri[0]};base64,{url}" + media_type, binary = _encode_image_binary(image, "png") + url = base64.b64encode(binary).decode("utf-8") + return f"data:{media_type};base64,{url}" except (IOError, FileNotFoundError): warnings.warn( f"Failed to read image {url}, with image_root set to {image_root}.", diff --git a/src/viser/_messages.py b/src/viser/_messages.py index 1e6d6517a..17db40451 100644 --- a/src/viser/_messages.py +++ b/src/viser/_messages.py @@ -153,7 +153,7 @@ class CameraFrustumMessage(Message): scale: float color: int image_media_type: Optional[Literal["image/jpeg", "image/png"]] - image_base64_data: Optional[str] + image_binary: Optional[bytes] @dataclasses.dataclass @@ -426,8 +426,8 @@ class BackgroundImageMessage(Message): """Message for rendering a background image.""" media_type: Literal["image/jpeg", "image/png"] - base64_rgb: str - base64_depth: Optional[str] + rgb_bytes: bytes + depth_bytes: Optional[bytes] @dataclasses.dataclass @@ -436,7 +436,7 @@ class ImageMessage(Message): name: str media_type: Literal["image/jpeg", "image/png"] - base64_data: str + data: bytes render_width: float render_height: float diff --git a/src/viser/_scene_api.py b/src/viser/_scene_api.py index 56fce9298..649abae94 100644 --- a/src/viser/_scene_api.py +++ b/src/viser/_scene_api.py @@ -1,6 +1,5 @@ from __future__ import annotations -import base64 import io import time import warnings @@ -73,11 +72,11 @@ def _encode_rgb(rgb: RgbTupleOrArray) -> int: return int(rgb_fixed[0] * (256**2) + rgb_fixed[1] * 256 + rgb_fixed[2]) -def _encode_image_base64( +def _encode_image_binary( image: onp.ndarray, format: Literal["png", "jpeg"], jpeg_quality: int | None = None, -) -> tuple[Literal["image/png", "image/jpeg"], str]: +) -> tuple[Literal["image/png", "image/jpeg"], bytes]: media_type: Literal["image/png", "image/jpeg"] image = _colors_to_uint8(image) with io.BytesIO() as data_buffer: @@ -94,10 +93,8 @@ def _encode_image_base64( ) else: assert_never(format) - - base64_data = base64.b64encode(data_buffer.getvalue()).decode("ascii") - - return media_type, base64_data + binary = data_buffer.getvalue() + return media_type, binary TVector = TypeVar("TVector", bound=tuple) @@ -444,12 +441,12 @@ def add_camera_frustum( """ if image is not None: - media_type, base64_data = _encode_image_base64( + media_type, binary = _encode_image_binary( image, format, jpeg_quality=jpeg_quality ) else: media_type = None - base64_data = None + binary = None self._websock_interface.queue_message( _messages.CameraFrustumMessage( @@ -460,7 +457,7 @@ def add_camera_frustum( # (255, 255, 255) => 0xffffff, etc color=_encode_rgb(color), image_media_type=media_type, - image_base64_data=base64_data, + image_binary=binary, ) ) return CameraFrustumHandle._make(self, name, wxyz, position, visible) @@ -1099,13 +1096,13 @@ def set_background_image( jpeg_quality: Quality of the jpeg image (if jpeg format is used). depth: Optional depth image to use to composite background with scene elements. """ - media_type, base64_data = _encode_image_base64( + media_type, rgb_bytes = _encode_image_binary( image, format, jpeg_quality=jpeg_quality ) # Encode depth if provided. We use a 3-channel PNG to represent a fixed point # depth at each pixel. - depth_base64data = None + depth_bytes = None if depth is not None: # Convert to fixed-point. # We'll support from 0 -> (2^24 - 1) / 100_000. @@ -1120,15 +1117,13 @@ def set_background_image( assert intdepth.shape == (*depth.shape[:2], 4) with io.BytesIO() as data_buffer: iio.imwrite(data_buffer, intdepth[:, :, :3], extension=".png") - depth_base64data = base64.b64encode(data_buffer.getvalue()).decode( - "ascii" - ) + depth_bytes = data_buffer.getvalue() self._websock_interface.queue_message( _messages.BackgroundImageMessage( media_type=media_type, - base64_rgb=base64_data, - base64_depth=depth_base64data, + rgb_bytes=rgb_bytes, + depth_bytes=depth_bytes, ) ) @@ -1162,14 +1157,14 @@ def add_image( Handle for manipulating scene node. """ - media_type, base64_data = _encode_image_base64( + media_type, binary = _encode_image_binary( image, format, jpeg_quality=jpeg_quality ) self._websock_interface.queue_message( _messages.ImageMessage( name=name, media_type=media_type, - base64_data=base64_data, + data=binary, render_width=render_width, render_height=render_height, ) diff --git a/src/viser/client/src/App.tsx b/src/viser/client/src/App.tsx index 9e9eb180d..e35cd3f09 100644 --- a/src/viser/client/src/App.tsx +++ b/src/viser/client/src/App.tsx @@ -464,11 +464,11 @@ function AdaptiveDpr() { const setDpr = useThree((state) => state.setDpr); return ( (refreshrate > 90 ? [80, 90] : [50, 60])} + step={0.1} + bounds={(refreshrate) => (refreshrate > 90 ? [40, 85] : [40, 55])} onChange={({ factor, fps, refreshrate }) => { const dpr = window.devicePixelRatio * (0.2 + 0.8 * factor); console.log( diff --git a/src/viser/client/src/MessageHandler.tsx b/src/viser/client/src/MessageHandler.tsx index 3d43408c0..0c10c0fe4 100644 --- a/src/viser/client/src/MessageHandler.tsx +++ b/src/viser/client/src/MessageHandler.tsx @@ -531,13 +531,18 @@ function useMessageHandler() { } // Add a camera frustum. case "CameraFrustumMessage": { - const texture = + let texture = undefined; + if ( message.image_media_type !== null && - message.image_base64_data !== null - ? new TextureLoader().load( - `data:${message.image_media_type};base64,${message.image_base64_data}`, - ) - : undefined; + message.image_binary !== null + ) { + const image_url = URL.createObjectURL( + new Blob([message.image_binary]), + ); + texture = new TextureLoader().load(image_url, () => + URL.revokeObjectURL(image_url), + ); + } addSceneNodeMakeParents( new SceneNode( @@ -709,34 +714,40 @@ function useMessageHandler() { } // Add a background image. case "BackgroundImageMessage": { - new TextureLoader().load( - `data:${message.media_type};base64,${message.base64_rgb}`, - (texture) => { - const oldBackgroundTexture = - viewer.backgroundMaterialRef.current!.uniforms.colorMap.value; - viewer.backgroundMaterialRef.current!.uniforms.colorMap.value = - texture; - if (isTexture(oldBackgroundTexture)) oldBackgroundTexture.dispose(); - - viewer.useGui.setState({ backgroundAvailable: true }); - }, + const rgb_url = URL.createObjectURL( + new Blob([message.rgb_bytes], { + type: message.media_type, + }), ); + new TextureLoader().load(rgb_url, (texture) => { + URL.revokeObjectURL(rgb_url); + const oldBackgroundTexture = + viewer.backgroundMaterialRef.current!.uniforms.colorMap.value; + viewer.backgroundMaterialRef.current!.uniforms.colorMap.value = + texture; + if (isTexture(oldBackgroundTexture)) oldBackgroundTexture.dispose(); + + viewer.useGui.setState({ backgroundAvailable: true }); + }); viewer.backgroundMaterialRef.current!.uniforms.enabled.value = true; viewer.backgroundMaterialRef.current!.uniforms.hasDepth.value = - message.base64_depth !== null; + message.depth_bytes !== null; - if (message.base64_depth !== null) { + if (message.depth_bytes !== null) { // If depth is available set the texture - new TextureLoader().load( - `data:image/png;base64,${message.base64_depth}`, - (texture) => { - const oldDepthTexture = - viewer.backgroundMaterialRef.current?.uniforms.depthMap.value; - viewer.backgroundMaterialRef.current!.uniforms.depthMap.value = - texture; - if (isTexture(oldDepthTexture)) oldDepthTexture.dispose(); - }, + const depth_url = URL.createObjectURL( + new Blob([message.depth_bytes], { + type: message.media_type, + }), ); + new TextureLoader().load(depth_url, (texture) => { + URL.revokeObjectURL(depth_url); + const oldDepthTexture = + viewer.backgroundMaterialRef.current?.uniforms.depthMap.value; + viewer.backgroundMaterialRef.current!.uniforms.depthMap.value = + texture; + if (isTexture(oldDepthTexture)) oldDepthTexture.dispose(); + }); } return; } @@ -831,8 +842,14 @@ function useMessageHandler() { // `addSceneNodeMakeParents` needs to be called immediately: it // overwrites position/wxyz attributes, and we don't want this to // happen after later messages are received. + const image_url = URL.createObjectURL( + new Blob([message.data], { + type: message.media_type, + }), + ); const texture = new TextureLoader().load( - `data:${message.media_type};base64,${message.base64_data}`, + image_url, + () => URL.revokeObjectURL(image_url), // Revoke URL on load. ); addSceneNodeMakeParents( new SceneNode( diff --git a/src/viser/client/src/WebsocketMessages.tsx b/src/viser/client/src/WebsocketMessages.tsx index 5155b6ab9..a09df19e4 100644 --- a/src/viser/client/src/WebsocketMessages.tsx +++ b/src/viser/client/src/WebsocketMessages.tsx @@ -61,7 +61,7 @@ export interface CameraFrustumMessage { scale: number; color: number; image_media_type: "image/jpeg" | "image/png" | null; - image_base64_data: string | null; + image_binary: Uint8Array | null; } /** GlTF Message * @@ -322,8 +322,8 @@ export interface TransformControlsUpdateMessage { export interface BackgroundImageMessage { type: "BackgroundImageMessage"; media_type: "image/jpeg" | "image/png"; - base64_rgb: string; - base64_depth: string | null; + rgb_bytes: Uint8Array; + depth_bytes: Uint8Array | null; } /** Message for rendering 2D images. * @@ -333,7 +333,7 @@ export interface ImageMessage { type: "ImageMessage"; name: string; media_type: "image/jpeg" | "image/png"; - base64_data: string; + data: Uint8Array; render_width: number; render_height: number; }