Merge pull request #173 from facebookresearch/bvh_fix

bvh now properly ignores non-root position channel frame data

animated_drawings/model/bvh.py

@@ -241,6 +241,19 @@ def _parse_skeleton(cls, lines: List[str]) -> BVH_Joint:
     def _process_frame_data(cls, skeleton: BVH_Joint, frames: List[List[float]]) -> Tuple[npt.NDArray[np.float32], npt.NDArray[np.float32]]:
         """ Given skeleton and frame data, return root position data and joint quaternion data, separately"""
 
+        def _get_frame_channel_order(joint: BVH_Joint, channels=[]):
+            channels.extend(joint.channel_order)
+            for child in [child for child in joint.get_children() if isinstance(child, BVH_Joint)]:
+                _get_frame_channel_order(child, channels)
+            return channels
+        channels = _get_frame_channel_order(skeleton)
+
+        # create a mask so we retain only joint rotations and root position
+        mask = np.array(list(map(lambda x: True if 'rotation' in x else False, channels)))
+        mask[:3] = True  # hack to make sure we keep root position
+
+        frames = np.array(frames, dtype=np.float32)[:, mask]
+
         # split root pose data and joint euler angle data
         pos_data, ea_rots = np.split(np.array(frames, dtype=np.float32), [3], axis=1)