Skinned mesh draft

src/viser/_message_api.py

@@ -799,6 +799,92 @@ def add_mesh(self, *args, **kwargs) -> MeshHandle:
         """Deprecated alias for `add_mesh_simple()`."""
         return self.add_mesh_simple(*args, **kwargs)
 
+    def add_mesh_skinned(
+        self,
+        name: str,
+        vertices: onp.ndarray,
+        faces: onp.ndarray,
+        bone_wxyzs: Tuple[Tuple[float, float, float, float], ...] | onp.ndarray,
+        bone_positions: Tuple[Tuple[float, float, float], ...] | onp.ndarray,
+        skin_weights: onp.ndarray,
+        color: RgbTupleOrArray = (90, 200, 255),
+        wireframe: bool = False,
+        opacity: Optional[float] = None,
+        material: Literal["standard", "toon3", "toon5"] = "standard",
+        flat_shading: bool = False,
+        side: Literal["front", "back", "double"] = "front",
+        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,
+    ) -> MeshHandle:
+        """Add a mesh to the scene.
+
+        Args:
+            name: A scene tree name. Names in the format of /parent/child can be used to
+                define a kinematic tree.
+            vertices: A numpy array of vertex positions. Should have shape (V, 3).
+            faces: A numpy array of faces, where each face is represented by indices of
+                vertices. Should have shape (F,)
+            bone_handles: Tuple of scene node handles. A bone will be attached to each.
+            skin_weights: A numpy array of skin weights. Should have shape (V, B) where B
+                is the number of bones.
+            color: Color of the mesh as an RGB tuple.
+            wireframe: Boolean indicating if the mesh should be rendered as a wireframe.
+            opacity: Opacity of the mesh. None means opaque.
+            material: Material type of the mesh ('standard', 'toon3', 'toon5').
+                This argument is ignored when wireframe=True.
+            flat_shading: Whether to do flat shading. This argument is ignored
+                when wireframe=True.
+            side: Side of the surface to render ('front', 'back', 'double').
+            wxyz: Quaternion rotation to parent frame from local frame (R_pl).
+            position: Translation from parent frame to local frame (t_pl).
+            visible: Whether or not this mesh is initially visible.
+
+        Returns:
+            Handle for manipulating scene node.
+        """
+        if wireframe and material != "standard":
+            warnings.warn(
+                f"Invalid combination of {wireframe=} and {material=}. Material argument will be ignored.",
+                stacklevel=2,
+            )
+        if wireframe and flat_shading:
+            warnings.warn(
+                f"Invalid combination of {wireframe=} and {flat_shading=}. Flat shading argument will be ignored.",
+                stacklevel=2,
+            )
+
+        assert skin_weights.shape == (vertices.shape[0], len(bone_handles))
+
+        # Take the four biggest indices.
+        top4_skin_indices = onp.argsort(skin_weights, axis=-1)[:, -4:]
+        top4_skin_weights = skin_weights[
+            onp.arange(vertices.shape[0])[:, None], top4_skin_indices
+        ]
+        assert (
+            top4_skin_weights.shape == top4_skin_indices.shape == (vertices.shape[0], 4)
+        )
+        self._queue(
+            _messages.MeshMessage(
+                name,
+                vertices.astype(onp.float32),
+                faces.astype(onp.uint32),
+                # (255, 255, 255) => 0xffffff, etc
+                color=_encode_rgb(color),
+                vertex_colors=None,
+                wireframe=wireframe,
+                opacity=opacity,
+                flat_shading=flat_shading,
+                side=side,
+                material=material,
+                bone_wxyzs=bone_wxyzs.astype(onp.float32),
+                bone_positions=bone_positions.astype(onp.float32),
+                skin_indices=top4_skin_indices.astype(onp.uint16),
+                skin_weights=top4_skin_weights.astype(onp.float32),
+            )
+        )
+        return MeshHandle._make(self, name, wxyz, position, visible)
+
     def add_mesh_simple(
         self,
         name: str,
@@ -861,6 +947,9 @@ def add_mesh_simple(
                 flat_shading=flat_shading,
                 side=side,
                 material=material,
+                bone_names=None,
+                skin_indices=None,
+                skin_weights=None,
             )
         )
         return MeshHandle._make(self, name, wxyz, position, visible)

src/viser/_messages.py

@@ -223,6 +223,13 @@ def __post_init__(self):
         assert self.colors.dtype == onp.uint8
 
 
+@dataclasses.dataclass
+class MeshBoneMessage(Message):
+    """Message for a bone of a skinned mesh."""
+
+    name: str
+
+
 @dataclasses.dataclass
 class MeshMessage(Message):
     """Mesh message.
@@ -248,6 +255,31 @@ def __post_init__(self):
         assert self.faces.shape[-1] == 3
 
 
+@dataclasses.dataclass
+class SkinnedMeshMessage(MeshMessage):
+    """Mesh message.
+
+    Vertices are internally canonicalized to float32, faces to uint32."""
+
+    bone_wxyzs: onpt.NDArray[onp.float32]
+    bone_positions: onpt.NDArray[onp.float32]
+    skin_indices: onpt.NDArray[onp.uint32]
+    skin_weights: onpt.NDArray[onp.float32]
+
+    def __post_init__(self):
+        # Check shapes.
+        assert self.vertices.shape[-1] == 3
+        assert self.faces.shape[-1] == 3
+        assert self.skin_weights is not None
+        assert (
+            self.skin_indices.shape
+            == self.skin_weights.shape
+            == (self.vertices.shape[0], 4)
+        )
+        assert self.bone_wxyzs.shape[-1] == 4
+        assert self.bone_positions.shape[-1] == 3
+
+
 @dataclasses.dataclass
 class TransformControlsMessage(Message):
     """Message for transform gizmos."""

src/viser/client/src/App.tsx

@@ -54,6 +54,8 @@ export type ViewerContextContents = {
   useSceneTree: UseSceneTree;
   useGui: UseGui;
   // Useful references.
+  // TODO: there's really no reason these all need to be their own ref objects.
+  // We could have just one ref to a global mutable struct.
   websocketRef: React.MutableRefObject<WebSocket | null>;
   canvasRef: React.MutableRefObject<HTMLCanvasElement | null>;
   sceneRef: React.MutableRefObject<THREE.Scene | null>;
@@ -75,6 +77,9 @@ export type ViewerContextContents = {
           overrideVisibility?: boolean; // Override from the GUI.
         };
   }>;
+  nodeRefFromName: React.MutableRefObject<{
+    [name: string]: undefined | THREE.Object3D;
+  }>;
   messageQueueRef: React.MutableRefObject<Message[]>;
   // Requested a render.
   getRenderRequestState: React.MutableRefObject<
@@ -137,6 +142,7 @@ function ViewerRoot() {
         })(),
       },
     }),
+    nodeRefFromName: React.useRef({}),
     messageQueueRef: React.useRef([]),
     getRenderRequestState: React.useRef("ready"),
     getRenderRequest: React.useRef(null),

src/viser/client/src/SceneTree.tsx

@@ -32,7 +32,8 @@ export class SceneNode<T extends THREE.Object3D = THREE.Object3D> {
      *
      * https://github.com/pmndrs/drei/issues/1323
      */
-    public readonly unmountWhenInvisible?: true,
+    public readonly unmountWhenInvisible?: boolean,
+    public readonly everyFrameCallback?: () => void,
   ) {
     this.children = [];
     this.clickable = false;
@@ -136,17 +137,20 @@ export function SceneNodeThreeObject(props: {
   const unmountWhenInvisible = viewer.useSceneTree(
     (state) => state.nodeFromName[props.name]?.unmountWhenInvisible,
   );
+  const everyFrameCallback = viewer.useSceneTree(
+    (state) => state.nodeFromName[props.name]?.everyFrameCallback,
+  );
   const [unmount, setUnmount] = React.useState(false);
   const clickable =
     viewer.useSceneTree((state) => state.nodeFromName[props.name]?.clickable) ??
     false;
   const [obj, setRef] = React.useState<THREE.Object3D | null>(null);
 
-  const dragInfo = React.useRef({
-    dragging: false,
-    startClientX: 0,
-    startClientY: 0,
-  });
+  // Update global registry of node objects.
+  // This is used for updating bone transforms in skinned meshes.
+  React.useEffect(() => {
+    if (obj !== null) viewer.nodeRefFromName.current[props.name] = obj;
+  }, [obj]);
 
   // Create object + children.
   //
@@ -206,6 +210,7 @@ export function SceneNodeThreeObject(props: {
   // although this shouldn't be a bottleneck.
   useFrame(() => {
     const attrs = viewer.nodeAttributesFromName.current[props.name];
+    everyFrameCallback && everyFrameCallback();
 
     // Unmount when invisible.
     // Examples: <Html /> components, PivotControls.
@@ -252,7 +257,12 @@ export function SceneNodeThreeObject(props: {
   });
 
   // Clean up when done.
-  React.useEffect(() => cleanup);
+  React.useEffect(() => {
+    return () => {
+      cleanup && cleanup();
+      delete viewer.nodeRefFromName.current[props.name];
+    };
+  });
 
   // Clicking logic.
   const sendClicksThrottled = makeThrottledMessageSender(
@@ -264,6 +274,12 @@ export function SceneNodeThreeObject(props: {
   const hoveredRef = React.useRef(false);
   if (!clickable && hovered) setHovered(false);
 
+  const dragInfo = React.useRef({
+    dragging: false,
+    startClientX: 0,
+    startClientY: 0,
+  });
+
   if (objNode === undefined || unmount) {
     return <>{children}</>;
   } else if (clickable) {
@@ -327,7 +343,6 @@ export function SceneNodeThreeObject(props: {
               });
             }}
             onPointerOver={(e) => {
-              console.log("over");
               if (!isDisplayed()) return;
               e.stopPropagation();
               setHovered(true);