diff --git a/process/skeleton_fitting.py b/process/skeleton_fitting.py
index 33721c5f2289ed7d60b5a6ec8580f8cfab2a0916..7a0177bf61f32dbb574fe2d45fda3783ef2226b4 100755
--- a/process/skeleton_fitting.py
+++ b/process/skeleton_fitting.py
@@ -142,6 +142,10 @@ class SkeletonFitting:
self.limbs_sets[limb_name] = self.limbs[limb_name].LimbsModuleSettings(init_yaml_path)
self.limbs_params[limb_name] = copy.deepcopy(self.limbs_sets[limb_name].params_init)
+ for param_name in self.body_params.keys():
+ self.limbs_params[limb_name]['right'][param_name] = self.body_params[param_name]
+ self.limbs_params[limb_name]['left'][param_name] = self.body_params[param_name]
+
self.skeleton2d, self.skeleton3d = None, None # type: Dict[str, any]
self.frame_rate = frame_rate # type: int
diff --git a/tests/test_skeleton_fitter.py b/tests/test_skeleton_fitter.py
index 098af14d65c89f0faaa8dd393c1002916cc500a3..70d232e7fe0adc617d6f38d7fd0510fdb801c0e9 100755
--- a/tests/test_skeleton_fitter.py
+++ b/tests/test_skeleton_fitter.py
@@ -1,4 +1,5 @@
import copy, os
+from matplotlib import pyplot as plt
from skeleton_fitter import optimiser, utils
from skeleton_fitter.modules.plotting import plot_limbs, plot_body_and_limbs
@@ -6,7 +7,7 @@ from skeleton_fitter.modules.plotting import plot_limbs, plot_body_and_limbs
from importlib import import_module
-def test_skeleton_fly(show_plots=False):
+def test_skeleton_fly(show_plots=True):
animal_names = ['fly', 'fly_geo', 'fly_hybrid']
for animal_name in animal_names:
@@ -37,10 +38,18 @@ def test_skeleton_fly(show_plots=False):
limbs_sets[limb_name] = limbs[limb_name].LimbsModuleSettings(init_yaml_path)
limbs_params[limb_name] = copy.deepcopy(limbs_sets[limb_name].params_init)
+ for param_name in body_params.keys():
+ limbs_params[limb_name]['right'][param_name] = body_params[param_name]
+ limbs_params[limb_name]['left'][param_name] = body_params[param_name]
+
limbs_params['wings']['right']['stroke_a'] = -55.0
limbs_params['wings']['right']['deviation_a'] = 10.0
limbs_params['wings']['right']['rotation_a'] = 60.0
+ limbs_params['wings']['left']['stroke_a'] = 35.0
+ limbs_params['wings']['left']['deviation_a'] = -10.0
+ limbs_params['wings']['left']['rotation_a'] = 30.0
+
wings_skeleton3d_1 = copy.deepcopy(limbs_sets['wings'].skeleton3d_init)
for side in limbs_sets['wings'].skeleton3d_init.keys():
hinge_label = '{0}_wing_hinge'.format(side)
@@ -56,17 +65,53 @@ def test_skeleton_fly(show_plots=False):
limbs_params['wings'][side], side)
if show_plots:
- plot_limbs(wings_skeleton3d_1)
- plot_body_and_limbs(body_skeleton3d_1, wings_skeleton3d_1)
+ # plot_limbs(wings_skeleton3d_1, wait_show_plot=True)
+ plot_body_and_limbs(body_skeleton3d_1, wings_skeleton3d_1, wait_show_plot=True)
+
+ # Estimate body and wings parameters
+ body_params1_est = body.estimate_params_from_skeleton3d(body_skeleton3d_1, body_sets.params_init)
+ limbs_params1_est = {}
+ for num_limb, limb_name in enumerate(animal_sets.limb_names):
+ limbs_params1_est[limb_name] = {}
+ for side in ['right', 'left']:
+ limbs_params1_est[limb_name][side] = (
+ limbs[limb_name].estimate_params_from_skeleton3d(wings_skeleton3d_1[side], body_params1_est,
+ limbs_sets[limb_name].params_init[side], side))
+
+ if show_plots:
+ body_skeleton3d_est = body.rotate_skeleton3d(body_sets.skeleton3d_init, body_params1_est)
+ body_skeleton3d_est = body.translate_skeleton3d(body_skeleton3d_est, body_params1_est)
+
+ wings_skeleton3d_est = copy.deepcopy(limbs_sets['wings'].skeleton3d_init)
+ for side in limbs_sets['wings'].skeleton3d_init.keys():
+ hinge_label = '{0}_wing_hinge'.format(side)
+ [limbs_params1_est['wings'][side]['x_hinge'], limbs_params1_est['wings'][side]['y_hinge'],
+ limbs_params1_est['wings'][side]['z_hinge']] = body_skeleton3d_est[hinge_label]
- # Estimate body and wings parameters
- # body_params1_est = estimate_body_parameters(body_skeleton3d_1, body_sets.params_init)
- # limbs_params1_est = {}
- # for side in ['right', 'left']:
- # limbs_params1_est[side] = estimate_wing_parameters(wings_skeleton3d_1[side], body_params1_est, limbs_sets[limb_name].params_init[side])
+ [limbs_params1_est['wings'][side]['x_com'], limbs_params1_est['wings'][side]['y_com'],
+ limbs_params1_est['wings'][side]['z_com']] = \
+ [body_params1_est['x_com'], body_params1_est['y_com'], body_params1_est['z_com']]
+ wings_skeleton3d_est[side] = (
+ limbs['wings'].rotate_and_translate_skeleton3d(limbs_sets['wings'].skeleton3d_init[side],
+ limbs_params1_est['wings'][side], side))
-def test_skeleton_fitter_fly(show_plots=False):
+ # plot_limbs(wings_skeleton3d_est, wait_show_plot=True)
+ plot_body_and_limbs(body_skeleton3d_est, wings_skeleton3d_est, wait_show_plot=True)
+
+ print('> differences between real value and estimate are:')
+ for label in ['yaw_a', 'pitch_a', 'roll_a']:
+ print(' {0} = {1} degrees'.format(label, body_params[label] - body_params1_est[label]))
+
+ for side in limbs_sets['wings'].skeleton3d_init.keys():
+ for label in ['stroke_a', 'deviation_a', 'rotation_a']:
+ print(' {0} {1} = {2} degrees'.format(label, side, limbs_params['wings'][side][label] -
+ limbs_params1_est['wings'][side][label]))
+
+ if show_plots: plt.show()
+
+
+def test_skeleton_fitter_fly():
max_rmse = 10
max_nb_iterations = 10000
@@ -106,6 +151,10 @@ def test_skeleton_fitter_fly(show_plots=False):
limbs_sets[limb_name] = limbs[limb_name].LimbsModuleSettings(init_yaml_path)
limbs_params[limb_name] = copy.deepcopy(limbs_sets[limb_name].params_init)
+ for param_name in body_params.keys():
+ limbs_params[limb_name]['right'][param_name] = body_params[param_name]
+ limbs_params[limb_name]['left'][param_name] = body_params[param_name]
+
# Parameters to generate fake 3d and 2d skeletons
body_params['yaw_a'] = 2.0
body_params['pitch_a'] = -1.0
@@ -173,7 +222,7 @@ def test_skeleton_fitter_fly(show_plots=False):
def main():
show_plots = True
test_skeleton_fly(show_plots)
- test_skeleton_fitter_fly(show_plots)
+ test_skeleton_fitter_fly()
if __name__ == "__main__":