Merge pull request #6 from BartMoyaers/master

Add joint_direct_children function

Author: rdeioris

README.md

@@ -72,3 +72,8 @@ Python module for parsing BVH (Biovision hierarchical data) mocap files
 >>> [str(node) for node in mocap.search('JOINT')]
 ['JOINT mixamorig:Spine', 'JOINT mixamorig:Spine1', 'JOINT mixamorig:Spine2', 'JOINT mixamorig:Neck', 'JOINT mixamorig:Head', 'JOINT mixamorig:HeadTop_End', 'JOINT mixamorig:LeftEye', 'JOINT mixamorig:RightEye', 'JOINT mixamorig:LeftShoulder', 'JOINT mixamorig:LeftArm', 'JOINT mixamorig:LeftForeArm', 'JOINT mixamorig:LeftHand', 'JOINT mixamorig:LeftHandMiddle1', 'JOINT mixamorig:LeftHandMiddle2', 'JOINT mixamorig:LeftHandMiddle3', 'JOINT mixamorig:LeftHandThumb1', 'JOINT mixamorig:LeftHandThumb2', 'JOINT mixamorig:LeftHandThumb3', 'JOINT mixamorig:LeftHandIndex1', 'JOINT mixamorig:LeftHandIndex2', 'JOINT mixamorig:LeftHandIndex3', 'JOINT mixamorig:LeftHandRing1', 'JOINT mixamorig:LeftHandRing2', 'JOINT mixamorig:LeftHandRing3', 'JOINT mixamorig:LeftHandPinky1', 'JOINT mixamorig:LeftHandPinky2', 'JOINT mixamorig:LeftHandPinky3', 'JOINT mixamorig:RightShoulder', 'JOINT mixamorig:RightArm', 'JOINT mixamorig:RightForeArm', 'JOINT mixamorig:RightHand', 'JOINT mixamorig:RightHandMiddle1', 'JOINT mixamorig:RightHandMiddle2', 'JOINT mixamorig:RightHandMiddle3', 'JOINT mixamorig:RightHandThumb1', 'JOINT mixamorig:RightHandThumb2', 'JOINT mixamorig:RightHandThumb3', 'JOINT mixamorig:RightHandIndex1', 'JOINT mixamorig:RightHandIndex2', 'JOINT mixamorig:RightHandIndex3', 'JOINT mixamorig:RightHandRing1', 'JOINT mixamorig:RightHandRing2', 'JOINT mixamorig:RightHandRing3', 'JOINT mixamorig:RightHandPinky1', 'JOINT mixamorig:RightHandPinky2', 'JOINT mixamorig:RightHandPinky3', 'JOINT mixamorig:RightUpLeg', 'JOINT mixamorig:RightLeg', 'JOINT mixamorig:RightFoot', 'JOINT mixamorig:RightToeBase', 'JOINT mixamorig:LeftUpLeg', 'JOINT mixamorig:LeftLeg', 'JOINT mixamorig:LeftFoot', 'JOINT mixamorig:LeftToeBase']
 ```
+#### Get joint's direct children
+```python
+>>> mocap.joint_direct_children('mixamorig:Hips')
+[JOINT mixamorig:Spine, JOINT mixamorig:RightUpLeg, JOINT mixamorig:LeftUpLeg]
+```

bvh.py

@@ -113,6 +113,10 @@ def iterate_joints(joint):
         iterate_joints(next(self.root.filter('ROOT')))
         return joints
 
+    def joint_direct_children(self, name):
+        joint = self.get_joint(name)
+        return [child for child in joint.filter('JOINT')]
+
     def get_joint_index(self, name):
         return self.get_joints().index(self.get_joint(name))
 

tests/test_bvh.py

@@ -173,6 +173,14 @@ def test_frames_multi_channels(self):
         rotations = mocap.frames_joint_channels('Head', ['Xrotation', 'Yrotation', 'Zrotation'])
         self.assertEqual(len(rotations), mocap.nframes)
 
-
+    def test_joint_children(self):
+        with open('tests/test_mocapbank.bvh') as f:
+            mocap = Bvh(f.read())
+        self.assertEqual(mocap.joint_direct_children('Chest')[0].name, 'Chest2')
+        self.assertEqual(mocap.joint_direct_children('Hips')[0].name, 'Chest')
+        self.assertEqual(mocap.joint_direct_children('Hips')[1].name, 'LeftHip')
+        self.assertEqual(mocap.joint_direct_children('Hips')[2].name, 'RightHip')
+        self.assertEqual(mocap.joint_direct_children('RightWrist'), [])
+        
 if __name__ == '__main__':
     unittest.main()