From 7448e7fcb0742cfb53a61b597621c1bd093b295e Mon Sep 17 00:00:00 2001
From: Brent Yi <yibrenth@gmail.com>
Date: Fri, 16 Feb 2024 16:30:25 -0800
Subject: [PATCH] Configurable origin radius for coordinate frames (#172)

* Configurable origin radius for coordinate frames

* Bump version
---
 pyproject.toml                              |  2 +-
 src/viser/_message_api.py                   |  7 ++++-
 src/viser/_messages.py                      |  7 ++---
 src/viser/client/src/ThreeAssets.tsx        | 29 +++++++++------------
 src/viser/client/src/WebsocketInterface.tsx |  1 +
 src/viser/client/src/WebsocketMessages.tsx  |  1 +
 6 files changed, 25 insertions(+), 22 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 32e0ffc19..7933c7582 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "viser"
-version = "0.1.23"
+version = "0.1.24"
 description = "3D visualization + Python"
 readme = "README.md"
 license = { text="MIT" }
diff --git a/src/viser/_message_api.py b/src/viser/_message_api.py
index 12efd8a5a..19c5edd35 100644
--- a/src/viser/_message_api.py
+++ b/src/viser/_message_api.py
@@ -550,6 +550,7 @@ def add_frame(
         show_axes: bool = True,
         axes_length: float = 0.5,
         axes_radius: float = 0.025,
+        origin_radius: float | None = None,
         wxyz: Tuple[float, float, float, float] | onp.ndarray = (1.0, 0.0, 0.0, 0.0),
         position: Tuple[float, float, float] | onp.ndarray = (0.0, 0.0, 0.0),
         visible: bool = True,
@@ -568,9 +569,10 @@ def add_frame(
         Args:
             name: A scene tree name. Names in the format of /parent/child can be used to
                 define a kinematic tree.
-            show_axes: Boolean to indicate whether to show the axes.
+            show_axes: Boolean to indicate whether to show the frame as a set of axes + origin sphere.
             axes_length: Length of each axis.
             axes_radius: Radius of each axis.
+            origin_radius: Radius of the origin sphere. If not set, defaults to `2 * axes_radius`.
             wxyz: Quaternion rotation to parent frame from local frame (R_pl).
             position: Translation to parent frame from local frame (t_pl).
             visible: Whether or not this scene node is initially visible.
@@ -578,12 +580,15 @@ def add_frame(
         Returns:
             Handle for manipulating scene node.
         """
+        if origin_radius is None:
+            origin_radius = axes_radius * 2
         self._queue(
             _messages.FrameMessage(
                 name=name,
                 show_axes=show_axes,
                 axes_length=axes_length,
                 axes_radius=axes_radius,
+                origin_radius=origin_radius,
             )
         )
         return FrameHandle._make(self, name, wxyz, position, visible)
diff --git a/src/viser/_messages.py b/src/viser/_messages.py
index 452658695..2b1bd6058 100644
--- a/src/viser/_messages.py
+++ b/src/viser/_messages.py
@@ -116,9 +116,10 @@ class FrameMessage(Message):
     """Coordinate frame message."""
 
     name: str
-    show_axes: bool = True
-    axes_length: float = 0.5
-    axes_radius: float = 0.025
+    show_axes: bool
+    axes_length: float
+    axes_radius: float
+    origin_radius: float
 
 
 @dataclasses.dataclass
diff --git a/src/viser/client/src/ThreeAssets.tsx b/src/viser/client/src/ThreeAssets.tsx
index d98b2b2c1..0d5fbe557 100644
--- a/src/viser/client/src/ThreeAssets.tsx
+++ b/src/viser/client/src/ThreeAssets.tsx
@@ -246,50 +246,45 @@ export const CoordinateFrame = React.forwardRef<
     showAxes?: boolean;
     axesLength?: number;
     axesRadius?: number;
+    originRadius?: number;
   }
 >(function CoordinateFrame(
   {
-    showAxes: show_axes = true,
-    axesLength: axes_length = 0.5,
-    axesRadius: axes_radius = 0.0125,
+    showAxes = true,
+    axesLength = 0.5,
+    axesRadius = 0.0125,
+    originRadius = undefined,
   },
   ref,
 ) {
+  originRadius = originRadius ?? axesRadius * 2;
   return (
     <group ref={ref}>
-      {show_axes && (
+      {showAxes && (
         <>
           <mesh
             geometry={originGeom}
             material={originMaterial}
-            scale={
-              new THREE.Vector3(
-                axes_radius * 2.5,
-                axes_radius * 2.5,
-                axes_radius * 2.5,
-              )
-            }
+            scale={new THREE.Vector3(originRadius, originRadius, originRadius)}
           >
             <OutlinesIfHovered />
           </mesh>
           <Instances limit={3}>
             <meshBasicMaterial />
-            <cylinderGeometry
-              args={[axes_radius, axes_radius, axes_length, 16]}
-            />
+            <cylinderGeometry args={[axesRadius, axesRadius, axesLength, 16]} />
             <Instance
               rotation={new THREE.Euler(0.0, 0.0, (3.0 * Math.PI) / 2.0)}
-              position={[0.5 * axes_length, 0.0, 0.0]}
+              position={[0.5 * axesLength, 0.0, 0.0]}
               color={0xcc0000}
             >
               <OutlinesIfHovered />
             </Instance>
-            <Instance position={[0.0, 0.5 * axes_length, 0.0]} color={0x00cc00}>
+            <Instance position={[0.0, 0.5 * axesLength, 0.0]} color={0x00cc00}>
               <OutlinesIfHovered />
             </Instance>
             <Instance
               rotation={new THREE.Euler(Math.PI / 2.0, 0.0, 0.0)}
-              position={[0.0, 0.0, 0.5 * axes_length]}
+              position={[0.0, 0.0, 0.5 * axesLength]}
               color={0x0000cc}
             >
               <OutlinesIfHovered />
diff --git a/src/viser/client/src/WebsocketInterface.tsx b/src/viser/client/src/WebsocketInterface.tsx
index 0ed6f0e41..408f93c64 100644
--- a/src/viser/client/src/WebsocketInterface.tsx
+++ b/src/viser/client/src/WebsocketInterface.tsx
@@ -147,6 +147,7 @@ function useMessageHandler() {
               showAxes={message.show_axes}
               axesLength={message.axes_length}
               axesRadius={message.axes_radius}
+              originRadius={message.origin_radius}
             />
           )),
         );
diff --git a/src/viser/client/src/WebsocketMessages.tsx b/src/viser/client/src/WebsocketMessages.tsx
index 08139e525..25a52bd41 100644
--- a/src/viser/client/src/WebsocketMessages.tsx
+++ b/src/viser/client/src/WebsocketMessages.tsx
@@ -71,6 +71,7 @@ export interface FrameMessage {
   show_axes: boolean;
   axes_length: number;
   axes_radius: number;
+  origin_radius: number;
 }
 /** Batched axes message.
  *