From f2ae8c1b35a4dcca9397cfa33054a583f4d81b46 Mon Sep 17 00:00:00 2001
From: brentyi <yibrenth@gmail.com>
Date: Fri, 8 Nov 2024 02:21:45 -0800
Subject: [PATCH] Add `origin_color` prop to frame objects

---
 src/viser/_messages.py                    | 2 ++
 src/viser/_scene_api.py                   | 2 ++
 src/viser/client/src/SceneTree.tsx        | 1 +
 src/viser/client/src/SceneTreeState.tsx   | 2 ++
 src/viser/client/src/ThreeAssets.tsx      | 5 +++--
 src/viser/client/src/WebsocketMessages.ts | 1 +
 6 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/src/viser/_messages.py b/src/viser/_messages.py
index 9922c0f2..4cca22e1 100644
--- a/src/viser/_messages.py
+++ b/src/viser/_messages.py
@@ -228,6 +228,8 @@ class FrameProps:
     """Radius of each axis. Synchronized automatically when assigned."""
     origin_radius: float
     """Radius of the origin sphere. Synchronized automatically when assigned."""
+    origin_color: Tuple[int, int, int]
+    """Color of the origin sphere as RGB integers. Synchronized automatically when assigned."""
 
 
 @dataclasses.dataclass
diff --git a/src/viser/_scene_api.py b/src/viser/_scene_api.py
index 2c847a82..fc7ebe8f 100644
--- a/src/viser/_scene_api.py
+++ b/src/viser/_scene_api.py
@@ -802,6 +802,7 @@ def add_frame(
         axes_length: float = 0.5,
         axes_radius: float = 0.025,
         origin_radius: float | None = None,
+        origin_color: RgbTupleOrArray = (236, 236, 0),
         wxyz: tuple[float, float, float, float] | np.ndarray = (1.0, 0.0, 0.0, 0.0),
         position: tuple[float, float, float] | np.ndarray = (0.0, 0.0, 0.0),
         visible: bool = True,
@@ -840,6 +841,7 @@ def add_frame(
                 axes_length=axes_length,
                 axes_radius=axes_radius,
                 origin_radius=origin_radius,
+                origin_color=_encode_rgb(origin_color),
             ),
         )
         return FrameHandle._make(self, message, name, wxyz, position, visible)
diff --git a/src/viser/client/src/SceneTree.tsx b/src/viser/client/src/SceneTree.tsx
index 91d10287..d5bbf3bf 100644
--- a/src/viser/client/src/SceneTree.tsx
+++ b/src/viser/client/src/SceneTree.tsx
@@ -147,6 +147,7 @@ function useObjectFactory(message: SceneNodeMessage | undefined): {
             axesLength={message.props.axes_length}
             axesRadius={message.props.axes_radius}
             originRadius={message.props.origin_radius}
+            originColor={rgbToInt(message.props.origin_color)}
           />
         ),
       };
diff --git a/src/viser/client/src/SceneTreeState.tsx b/src/viser/client/src/SceneTreeState.tsx
index dc81a0f4..cddbd070 100644
--- a/src/viser/client/src/SceneTreeState.tsx
+++ b/src/viser/client/src/SceneTreeState.tsx
@@ -37,6 +37,7 @@ export const rootNodeTemplate: SceneNode = {
       axes_length: 0.5,
       axes_radius: 0.0125,
       origin_radius: 0.025,
+      origin_color: [236, 236, 0],
     },
   },
   children: ["/WorldAxes"],
@@ -51,6 +52,7 @@ const worldAxesNodeTemplate: SceneNode = {
       axes_length: 0.5,
       axes_radius: 0.0125,
       origin_radius: 0.025,
+      origin_color: [236, 236, 0],
     },
   },
   children: [],
diff --git a/src/viser/client/src/ThreeAssets.tsx b/src/viser/client/src/ThreeAssets.tsx
index 36b49bae..49c74b96 100644
--- a/src/viser/client/src/ThreeAssets.tsx
+++ b/src/viser/client/src/ThreeAssets.tsx
@@ -54,7 +54,6 @@ function rgbToInt(rgb: [number, number, number]): number {
   return (rgb[0] << 16) | (rgb[1] << 8) | rgb[2];
 }
 const originGeom = new THREE.SphereGeometry(1.0);
-const originMaterial = new THREE.MeshBasicMaterial({ color: 0xecec00 });
 
 const PointCloudMaterial = /* @__PURE__ */ shaderMaterial(
   { scale: 1.0, point_ball_norm: 0.0 },
@@ -259,6 +258,7 @@ export const CoordinateFrame = React.forwardRef<
     axesLength?: number;
     axesRadius?: number;
     originRadius?: number;
+    originColor?: number;
   }
 >(function CoordinateFrame(
   {
@@ -266,6 +266,7 @@ export const CoordinateFrame = React.forwardRef<
     axesLength = 0.5,
     axesRadius = 0.0125,
     originRadius = undefined,
+    originColor = 0xecec00,
   },
   ref,
 ) {
@@ -276,9 +277,9 @@ export const CoordinateFrame = React.forwardRef<
         <>
           <mesh
             geometry={originGeom}
-            material={originMaterial}
             scale={new THREE.Vector3(originRadius, originRadius, originRadius)}
           >
+            <meshBasicMaterial color={originColor} />
             <OutlinesIfHovered />
           </mesh>
           <Instances limit={3}>
diff --git a/src/viser/client/src/WebsocketMessages.ts b/src/viser/client/src/WebsocketMessages.ts
index 66291459..10cde9b9 100644
--- a/src/viser/client/src/WebsocketMessages.ts
+++ b/src/viser/client/src/WebsocketMessages.ts
@@ -126,6 +126,7 @@ export interface FrameMessage {
     axes_length: number;
     axes_radius: number;
     origin_radius: number;
+    origin_color: [number, number, number];
   };
 }
 /** Batched axes message.