Added from_bvh method to skeleton

aitviewer/renderables/skeletons.py

@@ -20,6 +20,7 @@
 from aitviewer.renderables.spheres import Spheres
 from aitviewer.scene.material import Material
 from aitviewer.scene.node import Node
+from aitviewer.utils.bvh import Bvh
 
 
 class Skeletons(Node):
@@ -71,6 +72,40 @@ def __init__(
         )
         self._add_nodes(self.spheres, self.lines, show_in_hierarchy=False)
 
+    @classmethod
+    def from_bvh(cls, path: str, z_up=False, **kwargs):
+        """
+        Load an animated skeleton from a BVH (Biovision hierarchical data) mocap file.
+        :param path: path to the BVH file.
+        :param z_up: if True transform data from Z up to Y up.
+        :param kwargs: arguments forwarded to the Skeleton constructor.
+        """
+        # Parse BVH file
+        bvh = Bvh()
+        bvh.parse_string(open(path).read())
+
+        # Get positions for all frames.
+        positions, rotations = bvh.all_frame_poses()
+
+        # Recursively add connections.
+        connections = []
+        joints = list(bvh.joints.values())
+
+        def add_connections(node, index):
+            for c in node.children:
+                child_index = joints.index(c)
+                connections.append((index, child_index))
+                add_connections(c, child_index)
+
+        add_connections(bvh.root, 0)
+
+        # Transform to y up if data is z up.
+        rotation = kwargs.get("rotation", np.eye(3))
+        if z_up:
+            rotation = np.matmul(np.array([[1, 0, 0], [0, 0, 1], [0, -1, 0]]), rotation)
+
+        return cls(positions * 1e-2, np.array(connections), rotation=rotation)
+
     @property
     def joint_positions(self):
         return self._joint_positions

aitviewer/utils/bvh.py

@@ -0,0 +1,236 @@
+"""
+Biovision Hierarchy (BVH) reader for python with numpy pose output.
+
+Modified from:
+https://github.com/dabeschte/npybvh/tree/master
+
+
+MIT License
+
+Copyright (c) 2019 Mathias Parger
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+"""
+import re
+
+import numpy as np
+
+
+class BvhJoint:
+    def __init__(self, name, parent):
+        self.name = name
+        self.parent = parent
+        self.offset = np.zeros(3)
+        self.channels = []
+        self.children = []
+
+    def add_child(self, child):
+        self.children.append(child)
+
+    def __repr__(self):
+        return self.name
+
+    def position_animated(self):
+        return any([x.endswith("position") for x in self.channels])
+
+    def rotation_animated(self):
+        return any([x.endswith("rotation") for x in self.channels])
+
+
+class Bvh:
+    def __init__(self):
+        self.joints = {}
+        self.root = None
+        self.keyframes = None
+        self.frames = 0
+        self.fps = 0
+
+    def _parse_hierarchy(self, text):
+        lines = re.split("\\s*\\n+\\s*", text)
+
+        joint_stack = []
+
+        for line in lines:
+            words = re.split("\\s+", line)
+            instruction = words[0]
+
+            if instruction == "JOINT" or instruction == "ROOT":
+                parent = joint_stack[-1] if instruction == "JOINT" else None
+                joint = BvhJoint(words[1], parent)
+                self.joints[joint.name] = joint
+                if parent:
+                    parent.add_child(joint)
+                joint_stack.append(joint)
+                if instruction == "ROOT":
+                    self.root = joint
+            elif instruction == "CHANNELS":
+                for i in range(2, len(words)):
+                    joint_stack[-1].channels.append(words[i])
+            elif instruction == "OFFSET":
+                for i in range(1, len(words)):
+                    joint_stack[-1].offset[i - 1] = float(words[i])
+            elif instruction == "End":
+                parent = joint_stack[-1]
+                joint = BvhJoint(f"{parent.name}_end{len(parent.children)}", parent)
+                parent.add_child(joint)
+                joint_stack.append(joint)
+                self.joints[joint.name] = joint
+            elif instruction == "}":
+                joint_stack.pop()
+
+    def _add_pose_recursive(self, joint, offset, poses):
+        pose = joint.offset + offset
+        poses.append(pose)
+
+        for c in joint.children:
+            self._add_pose_recursive(c, pose, poses)
+
+    def plot_hierarchy(self):
+        import matplotlib.pyplot as plt
+
+        poses = []
+        self._add_pose_recursive(self.root, np.zeros(3), poses)
+        pos = np.array(poses)
+
+        fig = plt.figure()
+        ax = fig.add_subplot(111, projection="3d")
+        ax.scatter(pos[:, 0], pos[:, 2], pos[:, 1])
+        ax.set_xlim(-30, 30)
+        ax.set_ylim(-30, 30)
+        ax.set_zlim(-30, 30)
+        plt.show()
+
+    def parse_motion(self, text):
+        lines = re.split("\\s*\\n+\\s*", text)
+
+        frame = 0
+        for line in lines:
+            if line == "":
+                continue
+            words = re.split("\\s+", line)
+
+            if line.startswith("Frame Time:"):
+                self.fps = round(1 / float(words[2]))
+                continue
+            if line.startswith("Frames:"):
+                self.frames = int(words[1])
+                continue
+
+            if self.keyframes is None:
+                self.keyframes = np.empty((self.frames, len(words)), dtype=np.float32)
+
+            for angle_index in range(len(words)):
+                self.keyframes[frame, angle_index] = float(words[angle_index])
+
+            frame += 1
+
+    def parse_string(self, text):
+        hierarchy, motion = text.split("MOTION")
+        self._parse_hierarchy(hierarchy)
+        self.parse_motion(motion)
+
+    def _extract_rotation(self, frame_pose, index_offset, joint):
+        M_rotation = np.eye(3)
+        for channel in joint.channels:
+            if channel.endswith("position"):
+                continue
+            local_rotation = frame_pose[index_offset]
+            index_offset += 1
+            a = np.deg2rad(local_rotation)
+            s = np.sin(a)
+            c = np.cos(a)
+            if channel == "Xrotation":
+                M_channel = np.array([[1, 0, 0], [0, c, -s], [0, s, c]])
+            elif channel == "Yrotation":
+                M_channel = np.array([[c, 0, s], [0, 1, 0], [-s, 0, c]])
+            elif channel == "Zrotation":
+                M_channel = np.array([[c, -s, 0], [s, c, 0], [0, 0, 1]])
+            else:
+                raise Exception(f"Unknown channel {channel}")
+            M_rotation = M_rotation.dot(M_channel)
+
+        return M_rotation, index_offset
+
+    def _extract_position(self, joint, frame_pose, index_offset):
+        offset_position = np.zeros(3)
+        for channel in joint.channels:
+            if channel.endswith("rotation"):
+                continue
+            if channel == "Xposition":
+                offset_position[0] = frame_pose[index_offset]
+            elif channel == "Yposition":
+                offset_position[1] = frame_pose[index_offset]
+            elif channel == "Zposition":
+                offset_position[2] = frame_pose[index_offset]
+            else:
+                raise Exception(f"Unknown channel {channel}")
+            index_offset += 1
+
+        return offset_position, index_offset
+
+    def _recursive_apply_frame(self, joint, frame_pose, index_offset, p, r, M_parent, p_parent):
+        if joint.position_animated():
+            offset_position, index_offset = self._extract_position(joint, frame_pose, index_offset)
+        else:
+            offset_position = np.zeros(3)
+
+        if len(joint.channels) == 0:
+            joint_index = list(self.joints.values()).index(joint)
+            p[joint_index] = p_parent + M_parent.dot(joint.offset)
+            r[joint_index] = M_parent
+            return index_offset
+
+        if joint.rotation_animated():
+            M_rotation, index_offset = self._extract_rotation(frame_pose, index_offset, joint)
+        else:
+            M_rotation = np.eye(3)
+
+        M = M_parent.dot(M_rotation)
+        position = p_parent + M_parent.dot(joint.offset) + offset_position
+
+        rotation = M
+        joint_index = list(self.joints.values()).index(joint)
+        p[joint_index] = position
+        r[joint_index] = rotation
+
+        for c in joint.children:
+            index_offset = self._recursive_apply_frame(c, frame_pose, index_offset, p, r, M, position)
+
+        return index_offset
+
+    def frame_pose(self, frame):
+        p = np.empty((len(self.joints), 3))
+        r = np.empty((len(self.joints), 3, 3))
+        frame_pose = self.keyframes[frame]
+        M_parent = np.zeros((3, 3))
+        M_parent[0, 0] = 1
+        M_parent[1, 1] = 1
+        M_parent[2, 2] = 1
+        self._recursive_apply_frame(self.root, frame_pose, 0, p, r, M_parent, np.zeros(3))
+
+        return p, r
+
+    def all_frame_poses(self):
+        p = np.empty((self.frames, len(self.joints), 3))
+        r = np.empty((self.frames, len(self.joints), 3, 3))
+
+        for frame in range(len(self.keyframes)):
+            p[frame], r[frame] = self.frame_pose(frame)
+
+        return p, r