added optimizations

Author: Roberto De Ioris

bvh.py

@@ -149,6 +149,36 @@ def frame_joint_channel(self, frame_index, joint, channel, value=None):
             return value
         return float(self.frames[frame_index][joint_index + channel_index])
 
+    def frame_joint_channels(self, frame_index, joint, channels, value=None):
+        values = []
+        joint_index = self.get_joint_channels_index(joint)
+        for channel in channels:
+            channel_index = self.joint_channels(joint).index(channel)
+            if channel_index == -1 and value is not None:
+                values.append(value)
+            else:
+                values.append(
+                    float(
+                        self.frames[frame_index][joint_index + channel_index]
+                    )
+                )
+        return values
+
+    def frames_joint_channels(self, joint, channels, value=None):
+        all_frames = []
+        joint_index = self.get_joint_channels_index(joint)
+        for frame in self.frames:
+            values = []
+            for channel in channels:
+                channel_index = self.joint_channels(joint).index(channel)
+                if channel_index == -1 and value is not None:
+                    values.append(value)
+                else:
+                    values.append(
+                        float(frame[joint_index + channel_index]))
+            all_frames.append(values)
+        return all_frames
+
     def joint_parent(self, name):
         joint = self.get_joint(name)
         if joint.parent == self.root:

setup.py

@@ -3,7 +3,7 @@
 from distutils.core import setup
 
 setup(name='bvh',
-      version='0.2',
+      version='0.3',
       description='Python module for parsing BVH mocap files',
       author='20Tab S.r.l.',
       author_email='info@20tab.com',

tests/test_bvh.py

@@ -148,6 +148,18 @@ def test_joint_parent(self):
             mocap = Bvh(f.read())
         self.assertEqual(mocap.joint_parent('Chest').name, 'Hips')
 
+    def test_frame_joint_multi_channels(self):
+        with open('tests/test_mocapbank.bvh') as f:
+            mocap = Bvh(f.read())
+        rotation = mocap.frame_joint_channels(30, 'Head', ['Xrotation', 'Yrotation', 'Zrotation']) 
+        self.assertEqual(rotation, [1.77, 13.94, -7.42])
+
+    def test_frames_multi_channels(self):
+        with open('tests/test_mocapbank.bvh') as f:
+            mocap = Bvh(f.read())
+        rotations = mocap.frames_joint_channels('Head', ['Xrotation', 'Yrotation', 'Zrotation']) 
+        self.assertEqual(len(rotations), mocap.nframes)
+
 
 if __name__ == '__main__':
     unittest.main()