Add support for remote skeletons

Author: ramenguy99

aitviewer/remote/message.py

@@ -13,6 +13,7 @@ class Message(enum.Enum):
     ARROWS = 5
     RIGID_BODIES = 6
     SMPL = 10
+    SKELETONS = 11
 
     # Messages used to modify existing nodes on the remote viewer.
     DELETE = 100

aitviewer/remote/renderables/skeletons.py

@@ -0,0 +1,42 @@
+# Copyright (C) 2023  ETH Zurich, Manuel Kaufmann, Velko Vechev, Dario Mylonopoulos
+from ..message import Message
+from ..node import RemoteNode
+
+
+class RemoteSkeletons(RemoteNode):
+    MESSAGE_TYPE = Message.SKELETONS
+
+    def __init__(self, viewer, joint_positions, joint_connections, **kwargs):
+        """
+        This initializer takes a RemoteViewer object and all other arguments are forwarded
+        to the Spheres constructor on the remote Viewer.
+        See the Spheres class for more information about parameters.
+
+        :param viewer: a RemoteViewer object that will be used to send this node.
+        :param joint_positions: A np array of shape (F, J, 3) containing J joint positions over F many time steps.
+        :param joint_connections: The definition of the skeleton as a numpy array of shape (N_LINES, 2) where each row
+          defines one connection between joints. The max entry in this array must be < J.
+        """
+        super().__init__(
+            viewer,
+            joint_positions=joint_positions,
+            joint_connections=joint_connections,
+            **kwargs,
+        )
+
+    def add_frames(self, joint_positions):
+        """
+        Add frames to the remote Spheres node by adding new sphere positions.
+
+        :param joint_positions: A np array of shape (N, J, 3) or (J, 3) containing J joint positions.
+        """
+        return super().add_frames(joint_positions=joint_positions)
+
+    def update_frames(self, joint_positions, frames):
+        """
+        Update frames of the remote Spheres node by updating the sphere positions.
+
+        :param joint_positions: A np array of shape (N, J, 3) containing J joint positions.
+        :param frames: a list of integer frame indices of size N or a single integer frame index.
+        """
+        return super().update_frames(joint_positions=joint_positions, frames=frames)

aitviewer/renderables/skeletons.py

@@ -136,3 +136,19 @@ def color(self, color):
         self.material.color = color
         self.spheres.color = color
         self.lines.color = color
+
+    def update_frames(self, joint_positions, frames):
+        self.joint_positions[frames] = joint_positions
+        self.redraw()
+
+    def add_frames(self, joint_positions):
+        if len(joint_positions.shape) == 2:
+            joint_positions = joint_positions[np.newaxis]
+
+        self.joint_positions = np.append(self.joint_positions, joint_positions, axis=0)
+        self.spheres.add_frames(self.joint_positions)
+        self.lines.add_frames(self.joint_positions[:, self.skeleton].reshape(len(self), -1, 3))
+
+    def remove_frames(self, frames):
+        self.joint_positions = np.delete(self.joint_positions, frames, axis=0)
+        self.redraw()

aitviewer/server.py

@@ -13,6 +13,7 @@
 from aitviewer.renderables.lines import Lines
 from aitviewer.renderables.meshes import Meshes
 from aitviewer.renderables.rigid_bodies import RigidBodies
+from aitviewer.renderables.skeletons import Skeletons
 from aitviewer.renderables.smpl import SMPLSequence
 from aitviewer.renderables.spheres import Spheres
 from aitviewer.scene.node import Node
@@ -118,6 +119,9 @@ def add(client, remote_uid, args, kwargs, type):
         elif type == Message.RIGID_BODIES:
             add(client, remote_uid, args, kwargs, RigidBodies)
 
+        elif type == Message.SKELETONS:
+            add(client, remote_uid, args, kwargs, Skeletons)
+
         elif type == Message.SMPL:
             layer_arg_names = {"model_type", "gender", "num_betas"}
             layer_kwargs = {k: v for k, v in kwargs.items() if k in layer_arg_names}