From 78b55aa60b01ffd81338359c7e7cca92d4619fd1 Mon Sep 17 00:00:00 2001 From: Brent Yi Date: Mon, 29 Apr 2024 01:16:32 -0700 Subject: [PATCH] working again --- src/viser/_message_api.py | 39 ++++- src/viser/_messages.py | 36 ++++- src/viser/_scene_handles.py | 58 ++++++++ src/viser/client/src/App.tsx | 11 ++ src/viser/client/src/WebsocketInterface.tsx | 152 ++++++++++++-------- src/viser/client/src/WebsocketMessages.tsx | 30 +++- 6 files changed, 253 insertions(+), 73 deletions(-) diff --git a/src/viser/_message_api.py b/src/viser/_message_api.py index 52710f61e..31c5faea9 100644 --- a/src/viser/_message_api.py +++ b/src/viser/_message_api.py @@ -40,6 +40,8 @@ from . import transforms as tf from ._scene_handles import ( BatchedAxesHandle, + BoneHandle, + BoneState, CameraFrustumHandle, FrameHandle, GlbHandle, @@ -858,7 +860,8 @@ def add_mesh_skinned( stacklevel=2, ) - assert skin_weights.shape == (vertices.shape[0], len(bone_wxyzs)) + num_bones = len(bone_wxyzs) + assert skin_weights.shape == (vertices.shape[0], num_bones) # Take the four biggest indices. top4_skin_indices = onp.argsort(skin_weights, axis=-1)[:, -4:] @@ -871,6 +874,8 @@ def add_mesh_skinned( bone_wxyzs = onp.asarray(bone_wxyzs) bone_positions = onp.asarray(bone_positions) + assert bone_wxyzs.shape == (num_bones, 4) + assert bone_positions.shape == (num_bones, 3) self._queue( _messages.SkinnedMeshMessage( name, @@ -884,13 +889,39 @@ def add_mesh_skinned( flat_shading=flat_shading, side=side, material=material, - bone_wxyzs=bone_wxyzs.astype(onp.float32), - bone_positions=bone_positions.astype(onp.float32), + bone_wxyzs=tuple( + ( + float(wxyz[0]), + float(wxyz[1]), + float(wxyz[2]), + float(wxyz[3]), + ) + for wxyz in bone_wxyzs.astype(onp.float32) + ), + bone_positions=tuple( + (float(xyz[0]), float(xyz[1]), float(xyz[2])) + for xyz in bone_positions.astype(onp.float32) + ), skin_indices=top4_skin_indices.astype(onp.uint16), skin_weights=top4_skin_weights.astype(onp.float32), ) ) - return SkinnedMeshHandle._make(self, name, wxyz, position, visible) + handle = MeshHandle._make(self, name, wxyz, position, visible) + return SkinnedMeshHandle( + handle._impl, + bones=tuple( + BoneHandle( + _impl=BoneState( + name=name, + api=self, + bone_index=i, + wxyz=bone_wxyzs[i], + position=bone_positions[i], + ) + ) + for i in range(num_bones) + ), + ) def add_mesh_simple( self, diff --git a/src/viser/_messages.py b/src/viser/_messages.py index eabac1744..efc17f19b 100644 --- a/src/viser/_messages.py +++ b/src/viser/_messages.py @@ -261,8 +261,8 @@ class SkinnedMeshMessage(MeshMessage): Vertices are internally canonicalized to float32, faces to uint32.""" - bone_wxyzs: onpt.NDArray[onp.float32] - bone_positions: onpt.NDArray[onp.float32] + bone_wxyzs: Tuple[Tuple[float, float, float, float], ...] + bone_positions: Tuple[Tuple[float, float, float], ...] skin_indices: onpt.NDArray[onp.uint32] skin_weights: onpt.NDArray[onp.float32] @@ -276,8 +276,36 @@ def __post_init__(self): == 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 SetBoneOrientationMessage(Message): + """Server -> client message to set a skinned mesh bone's orientation. + + As with all other messages, transforms take the `T_parent_local` convention.""" + + name: str + bone_index: int + wxyz: Tuple[float, float, float, float] + + @override + def redundancy_key(self) -> str: + return type(self).__name__ + "-" + self.name + "-" + str(self.bone_index) + + +@dataclasses.dataclass +class SetBonePositionMessage(Message): + """Server -> client message to set a skinned mesh bone's position. + + As with all other messages, transforms take the `T_parent_local` convention.""" + + name: str + bone_index: int + position: Tuple[float, float, float] + + @override + def redundancy_key(self) -> str: + return type(self).__name__ + "-" + self.name + "-" + str(self.bone_index) @dataclasses.dataclass diff --git a/src/viser/_scene_handles.py b/src/viser/_scene_handles.py index c557d29fd..bcb25776b 100644 --- a/src/viser/_scene_handles.py +++ b/src/viser/_scene_handles.py @@ -218,6 +218,64 @@ class MeshHandle(_ClickableSceneNodeHandle): class SkinnedMeshHandle(_ClickableSceneNodeHandle): """Handle for skinned mesh objects.""" + bones: Tuple[BoneHandle, ...] + """Bones of the skinned mesh. These handles can be used for reading and + writing poses, which are defined relative to the mesh root.""" + + +@dataclasses.dataclass +class BoneState: + name: str + api: MessageApi + bone_index: int + wxyz: onp.ndarray + position: onp.ndarray + + +@dataclasses.dataclass +class BoneHandle: + """Handle for reading and writing the poses of bones in a skinned mesh.""" + + _impl: BoneState + + @property + def wxyz(self) -> onp.ndarray: + """Orientation of the bone. This is the quaternion representation of the R + in `p_parent = [R | t] p_local`. Synchronized to clients automatically when assigned. + """ + return self._impl.wxyz + + @wxyz.setter + def wxyz(self, wxyz: Tuple[float, float, float, float] | onp.ndarray) -> None: + from ._message_api import cast_vector + + wxyz_cast = cast_vector(wxyz, 4) + self._impl.wxyz = onp.asarray(wxyz) + self._impl.api._queue( + _messages.SetBoneOrientationMessage( + self._impl.name, self._impl.bone_index, wxyz_cast + ) + ) + + @property + def position(self) -> onp.ndarray: + """Position of the bone. This is equivalent to the t in + `p_parent = [R | t] p_local`. Synchronized to clients automatically when assigned. + """ + return self._impl.position + + @position.setter + def position(self, position: Tuple[float, float, float] | onp.ndarray) -> None: + from ._message_api import cast_vector + + position_cast = cast_vector(position, 3) + self._impl.position = onp.asarray(position) + self._impl.api._queue( + _messages.SetBonePositionMessage( + self._impl.name, self._impl.bone_index, position_cast + ) + ) + @dataclasses.dataclass class GlbHandle(_ClickableSceneNodeHandle): diff --git a/src/viser/client/src/App.tsx b/src/viser/client/src/App.tsx index 0552a1275..bbe185866 100644 --- a/src/viser/client/src/App.tsx +++ b/src/viser/client/src/App.tsx @@ -94,6 +94,16 @@ export type ViewerContextContents = { }>; // 2D canvas for drawing -- can be used to give feedback on cursor movement, or more. canvas2dRef: React.MutableRefObject; + // Poses for bones in skinned meshes. + skinnedMeshState: React.MutableRefObject<{ + [name: string]: { + initialized: boolean; + poses: { + wxyz: [number, number, number, number]; + position: [number, number, number]; + }[]; + }; + }>; }; export const ViewerContext = React.createContext( null, @@ -152,6 +162,7 @@ function ViewerRoot() { listening: false, }), canvas2dRef: React.useRef(null), + skinnedMeshState: React.useRef({}), }; return ( diff --git a/src/viser/client/src/WebsocketInterface.tsx b/src/viser/client/src/WebsocketInterface.tsx index d3f4277ba..f59712327 100644 --- a/src/viser/client/src/WebsocketInterface.tsx +++ b/src/viser/client/src/WebsocketInterface.tsx @@ -227,16 +227,20 @@ function useMessageHandler() { message.plane == "xz" ? new THREE.Euler(0.0, 0.0, 0.0) : message.plane == "xy" - ? new THREE.Euler(Math.PI / 2.0, 0.0, 0.0) - : message.plane == "yx" - ? new THREE.Euler(0.0, Math.PI / 2.0, Math.PI / 2.0) - : message.plane == "yz" - ? new THREE.Euler(0.0, 0.0, Math.PI / 2.0) - : message.plane == "zx" - ? new THREE.Euler(0.0, Math.PI / 2.0, 0.0) - : message.plane == "zy" - ? new THREE.Euler(-Math.PI / 2.0, 0.0, -Math.PI / 2.0) - : undefined + ? new THREE.Euler(Math.PI / 2.0, 0.0, 0.0) + : message.plane == "yx" + ? new THREE.Euler(0.0, Math.PI / 2.0, Math.PI / 2.0) + : message.plane == "yz" + ? new THREE.Euler(0.0, 0.0, Math.PI / 2.0) + : message.plane == "zx" + ? new THREE.Euler(0.0, Math.PI / 2.0, 0.0) + : message.plane == "zy" + ? new THREE.Euler( + -Math.PI / 2.0, + 0.0, + -Math.PI / 2.0, + ) + : undefined } /> @@ -324,16 +328,16 @@ function useMessageHandler() { message.material == "standard" || message.wireframe ? new THREE.MeshStandardMaterial(standardArgs) : message.material == "toon3" - ? new THREE.MeshToonMaterial({ - gradientMap: generateGradientMap(3), - ...standardArgs, - }) - : message.material == "toon5" - ? new THREE.MeshToonMaterial({ - gradientMap: generateGradientMap(5), - ...standardArgs, - }) - : assertUnreachable(message.material); + ? new THREE.MeshToonMaterial({ + gradientMap: generateGradientMap(3), + ...standardArgs, + }) + : message.material == "toon5" + ? new THREE.MeshToonMaterial({ + gradientMap: generateGradientMap(5), + ...standardArgs, + }) + : assertUnreachable(message.material); geometry.setAttribute( "position", new THREE.Float32BufferAttribute( @@ -388,47 +392,41 @@ function useMessageHandler() { cleanupMesh, ), ); - else if (message.type === "SkinMeshMessage") { - const getT_world_local: (name: string) => THREE.Matrix4 = ( - name: string, - ) => { - const T_current_local = new THREE.Matrix4().identity(); - const T_parent_current = new THREE.Matrix4().identity(); - let done = false; - while (!done) { - const attrs = viewer.nodeAttributesFromName.current[name]; - let wxyz = attrs?.wxyz; - if (wxyz === undefined) wxyz = [1, 0, 0, 0]; - T_parent_current.makeRotationFromQuaternion( - new THREE.Quaternion(wxyz[1], wxyz[2], wxyz[3], wxyz[0]), - ); - let position = attrs?.position; - if (position === undefined) position = [0, 0, 0]; - T_parent_current.setPosition( - new THREE.Vector3(position[0], position[1], position[2]), - ); - - T_current_local.premultiply(T_parent_current); - if (name === "") break; - name = name.split("/").slice(0, -1).join("/"); - console.log(name); - } - return T_current_local; - }; + else if (message.type === "SkinnedMeshMessage") { // Skinned mesh. const bones: THREE.Bone[] = []; - for (let i = 0; i < message.bone_names!.length; i++) { + for (let i = 0; i < message.bone_wxyzs!.length; i++) { bones.push(new THREE.Bone()); } + + const xyzw_quat = new THREE.Quaternion(); + const boneInverses: THREE.Matrix4[] = []; + viewer.skinnedMeshState.current[message.name] = { + initialized: false, + poses: [], + }; bones.forEach((bone, i) => { - scene.add(bone); - bone.matrix.copy(getT_world_local(message.bone_names![i])); - bone.matrixWorld.copy(getT_world_local(message.bone_names![i])); - // We'll manage the bone matrices manually. + const wxyz = message.bone_wxyzs[i]; + const position = message.bone_positions[i]; + xyzw_quat.set(wxyz[1], wxyz[2], wxyz[3], wxyz[0]); + + const boneInverse = new THREE.Matrix4(); + boneInverse.makeRotationFromQuaternion(xyzw_quat); + boneInverse.setPosition(position[0], position[1], position[2]); + boneInverse.invert(); + boneInverses.push(boneInverse); + + bone.quaternion.copy(xyzw_quat); + bone.position.set(position[0], position[1], position[2]); bone.matrixAutoUpdate = false; bone.matrixWorldAutoUpdate = false; + + viewer.skinnedMeshState.current[message.name].poses.push({ + wxyz: wxyz, + position: position, + }); }); - const skeleton = new THREE.Skeleton(bones); + const skeleton = new THREE.Skeleton(bones, boneInverses); geometry.setAttribute( "skinIndex", @@ -456,6 +454,7 @@ function useMessageHandler() { 4, ), ); + addSceneNodeMakeParents( new SceneNode( message.name, @@ -481,25 +480,52 @@ function useMessageHandler() { false, // everyFrameCallback: update bone transforms. () => { + const parentNode = viewer.nodeRefFromName.current[message.name]; + if (parentNode === undefined) return; + + const state = viewer.skinnedMeshState.current[message.name]; bones.forEach((bone, i) => { - const nodeRef = - viewer.nodeRefFromName.current[message.bone_names![i]]; - if (nodeRef !== undefined) { - // Our bone objects are placed in the scene root! - // bone.matrix.copy(nodeRef?.matrixWorld); - // bone.matrixWorld.copy(nodeRef?.matrixWorld); - bone.matrix.copy(getT_world_local(message.bone_names![i])); - bone.matrixWorld.copy( - getT_world_local(message.bone_names![i]), - ); + if (!state.initialized) { + parentNode.add(bone); } + const wxyz = state.initialized + ? state.poses[i].wxyz + : message.bone_wxyzs[i]; + const position = state.initialized + ? state.poses[i].position + : message.bone_positions[i]; + + xyzw_quat.set(wxyz[1], wxyz[2], wxyz[3], wxyz[0]); + bone.matrix.makeRotationFromQuaternion(xyzw_quat); + bone.matrix.setPosition( + position[0], + position[1], + position[2], + ); + bone.updateMatrixWorld(); }); + + if (!state.initialized) { + state.initialized = true; + } }, ), ); } return; } + // Set the bone poses. + case "SetBoneOrientationMessage": { + const bonePoses = viewer.skinnedMeshState.current; + bonePoses[message.name].poses[message.bone_index].wxyz = message.wxyz; + break; + } + case "SetBonePositionMessage": { + const bonePoses = viewer.skinnedMeshState.current; + bonePoses[message.name].poses[message.bone_index].position = + message.position; + break; + } // Add a camera frustum. case "CameraFrustumMessage": { const texture = diff --git a/src/viser/client/src/WebsocketMessages.tsx b/src/viser/client/src/WebsocketMessages.tsx index 275657926..e6b1b9e00 100644 --- a/src/viser/client/src/WebsocketMessages.tsx +++ b/src/viser/client/src/WebsocketMessages.tsx @@ -190,11 +190,35 @@ export interface SkinnedMeshMessage { flat_shading: boolean; side: "front" | "back" | "double"; material: "standard" | "toon3" | "toon5"; - bone_wxyzs: Uint8Array; - bone_positions: Uint8Array; + bone_wxyzs: [number, number, number, number][]; + bone_positions: [number, number, number][]; skin_indices: Uint8Array; skin_weights: Uint8Array; } +/** Server -> client message to set a skinned mesh bone's orientation. + * + * As with all other messages, transforms take the `T_parent_local` convention. + * + * (automatically generated) + */ +export interface SetBoneOrientationMessage { + type: "SetBoneOrientationMessage"; + name: string; + bone_index: number; + wxyz: [number, number, number, number]; +} +/** Server -> client message to set a skinned mesh bone's position. + * + * As with all other messages, transforms take the `T_parent_local` convention. + * + * (automatically generated) + */ +export interface SetBonePositionMessage { + type: "SetBonePositionMessage"; + name: string; + bone_index: number; + position: [number, number, number]; +} /** Message for transform gizmos. * * (automatically generated) @@ -860,6 +884,8 @@ export type Message = | MeshBoneMessage | MeshMessage | SkinnedMeshMessage + | SetBoneOrientationMessage + | SetBonePositionMessage | TransformControlsMessage | SetCameraPositionMessage | SetCameraUpDirectionMessage