diff --git a/README.md b/README.md index 89194ff1940517a80c16d474429fd6d6caea12a3..29c287fa65343164ace96478c1ebf0290629ad1c 100755 --- a/README.md +++ b/README.md @@ -105,7 +105,7 @@ Will stitch all views together for each frame. By stitching all views of the sam ### Select or define an animal skeleton An animal skeleton is defined as one body module + one or several (pair(s) of) limbs module(s). -For example, a [fly skeleton](/skeleton_fitter/animals/fly.py) can consist of an [insect body](/skeleton_fitter/modules/bodies/insect_body_slim.py) (three jointed segments for the head, thorax and abdomen + one segment between the two wing hinges) and two [wings](/skeleton_fitter/modules/limbs/insect_wings_flat.py) (here defined as flat). Each limb module consist of symmetrical limbs (one left and one right limb). +For example, a [fly skeleton](/skeleton_fitter/animals/fly.py) can consist of an [insect body](/skeleton_fitter/modules/bodies/insect_body_3segments_4wings.py) (three jointed segments for the head, thorax and abdomen + one segment between the two wing hinges) and two [wings](/skeleton_fitter/modules/limbs/insect_wings_flat.py) (here defined as flat). Each limb module consist of symmetrical limbs (one left and one right limb). You can select an already defined animal skeleton (see [/skeleton_fitter/animals](/skeleton_fitter/animals)). diff --git a/skeleton_fitter/animals/butterfly.py b/skeleton_fitter/animals/butterfly.py new file mode 100644 index 0000000000000000000000000000000000000000..4e625edd6d93c4a8be12fbfd1bfa009124608928 --- /dev/null +++ b/skeleton_fitter/animals/butterfly.py @@ -0,0 +1,24 @@ +from typing import List + + +class AnimalSettings: + """ Class for keeping track of the settings of a dragonfly skeleton. + + An animal skeleton consist of multiple modules (one body + various number of pairs of limbs (right and left)) + for example: a dragonfly consist of one body: insect_body_3segments_2wings, + and multiple limbs: two pairs of insect_wings_flat + """ + + body_module_name = 'insect_body_1segment_4wings' # type: str + limb_module_names = ['insect_wings_flat_with_apex', 'insect_wings_flat_with_apex'] # type: List[str] + limb_names = ['forewings', 'hindwings'] # type: str + + assert len(limb_module_names) == len(limb_names) + + @staticmethod + def from_dict(dictionary): + animal_sets = AnimalSettings() + for key, value in dictionary.items(): + setattr(animal_sets, key, value) + + return animal_sets diff --git a/skeleton_fitter/animals/dragonfly.py b/skeleton_fitter/animals/dragonfly.py index b0f3c02ba3addd9570af3471c9569bfdbeed13df..153c67eedb1accb55967ddb959b918e29201d9ce 100644 --- a/skeleton_fitter/animals/dragonfly.py +++ b/skeleton_fitter/animals/dragonfly.py @@ -5,11 +5,11 @@ class AnimalSettings: """ Class for keeping track of the settings of a dragonfly skeleton. An animal skeleton consist of multiple modules (one body + various number of pairs of limbs (right and left)) - for example: a dragonfly consist of one body: insect_body_slim, + for example: a dragonfly consist of one body: insect_body_3segments_2wings, and multiple limbs: two pairs of insect_wings_flat """ - body_module_name = 'insect_body_slim' # type: str + body_module_name = 'insect_body_3segments_4wings' # type: str limb_module_names = ['insect_wings_flat', 'insect_wings_flat'] # type: List[str] limb_names = ['forewings', 'hindwings'] # type: str diff --git a/skeleton_fitter/animals/fly.py b/skeleton_fitter/animals/fly.py index 90899063eba3d8a21ec72157df83d8926e00c405..13f6806abebeb00662191b3f84255f6040c501e4 100755 --- a/skeleton_fitter/animals/fly.py +++ b/skeleton_fitter/animals/fly.py @@ -5,11 +5,11 @@ class AnimalSettings: """ Class for keeping track of the settings of a fly skeleton. An animal skeleton consist of multiple modules (one body + various number of pairs of limbs (right and left)) - for example: a fly consist of one body: insect_body_slim, + for example: a fly consist of one body: insect_body_3segments_2wings, and multiple limbs: one pair of insect_wings_flat (+ three pairs of insect_legs) """ - body_module_name = 'insect_body_slim' # type: str + body_module_name = 'insect_body_3segments_2wings' # type: str limb_module_names = ['insect_wings_flat'] # type: List[str] limb_names = ['wings'] # type: str diff --git a/skeleton_fitter/animals/fly_geo.py b/skeleton_fitter/animals/fly_geo.py index fab79ce0e4842bc4d35496cebc21404c83cfe418..93021affb4d53225ea078820cff26f3e96d410c3 100755 --- a/skeleton_fitter/animals/fly_geo.py +++ b/skeleton_fitter/animals/fly_geo.py @@ -6,10 +6,10 @@ class AnimalSettings: Will be using the full geometry of the wing instead of only a few discrete points. An animal skeleton consist of multiple modules (one body + various number of pairs of limbs (right and left)) - for example: a fly consist of one body: insect_body_slim and one pair of insect_wings_flat_geo + for example: a fly consist of one body: insect_body_3segments_2wings and one pair of insect_wings_flat_geo """ - body_module_name = 'insect_body_slim' # type: str + body_module_name = 'insect_body_3segments_2wings' # type: str limb_module_names = ['insect_wings_flat_geo'] # type: List[str] limb_names = ['wings'] # type: str diff --git a/skeleton_fitter/animals/fly_hybrid.py b/skeleton_fitter/animals/fly_hybrid.py index 8dd66b566817ec370e38636a71830496ccd28068..aa538ac96ebe4c0a01e76af980224689bd4ee980 100755 --- a/skeleton_fitter/animals/fly_hybrid.py +++ b/skeleton_fitter/animals/fly_hybrid.py @@ -6,12 +6,12 @@ class AnimalSettings: This skeleton is using a hybrid body (which use info on both the body and the wings) An animal skeleton consist of multiple modules (one body + various number of pairs of limbs (right and left)). - Here a hybrid fly consist of one body: insect_hybrid_body_slim, and two limbs: one pair of insect_wings_flat + Here a hybrid fly consist of one body: insect_body_3segments_2wings_hybrid, and two limbs: one pair of insect_wings_flat """ use_hybrid_body = True # type: bool - body_module_name = 'insect_hybrid_body_slim' # type: str + body_module_name = 'insect_body_3segments_2wings_hybrid' # type: str limb_module_names = ['insect_wings_flat'] # type: List[str] limb_names = ['wings'] # type: str diff --git a/skeleton_fitter/modules/bodies/insect_body_1segment_4wings.py b/skeleton_fitter/modules/bodies/insect_body_1segment_4wings.py new file mode 100644 index 0000000000000000000000000000000000000000..dd63ebe1d8092cff09d096fc523f523282319ac9 --- /dev/null +++ b/skeleton_fitter/modules/bodies/insect_body_1segment_4wings.py @@ -0,0 +1,371 @@ +import copy, yaml +import numpy as np + +from dataclasses import dataclass + +from scipy.spatial.transform import Rotation + +from typing import Dict, List + +from skeleton_fitter.utils import reproject_skeleton3d_to2d, length_from_vector +from camera.utils import get_rotation_angle_btw_vectors + + +def init_params(init_dict: Dict[str, any]) -> Dict[str, float]: + """ + Load initial kinematic parameters + + Args: + init_dict: Dictionary with initial values of the skeleton coordinates/parameters/bounds + + Returns: + body_params: + """ + + param_names = ['default_pitch_a', 'yaw_a', 'pitch_a', 'roll_a', 'body_l', 'body_r', + 'ratio_com_body', 'ratio_body', + 'x_com', 'y_com', 'z_com'] + + body_params = dict() + for param_name in param_names: + body_params[param_name] = init_dict['init_params']['body'][param_name] + + return body_params + + +def init_bounds(init_dict: Dict[str, any]) -> Dict[str, float]: + """ + Load initial bounds for kinematic parameters + + Args: + init_dict: Dictionary with initial values of the skeleton coordinates/parameters/bounds + + Returns: + body_bounds: + """ + + param_names = ['default_pitch_a', 'yaw_a', 'pitch_a', 'roll_a', 'body_l', 'body_r', + 'ratio_com_body', 'ratio_body', + 'x_com', 'y_com', 'z_com'] + + body_bounds = dict() + for param_name in param_names: + body_bounds[param_name] = tuple(init_dict['init_bounds']['body'][param_name]) + + return body_bounds + + +def init_skeleton3d(init_dict: Dict[str, any], body_params: Dict[str, float]) -> Dict[str, List[float]]: + """ + Initialize body skeleton + + Args: + init_dict: Dictionary with initial values of the skeleton coordinates/parameters/bounds + body_params: + + Returns: + body_skeleton3d: + """ + + # Initialise a body with 3 segments (proboscis-head, torso, abdomen) along x axis + labels = ['head_tip', 'thorax_middle', 'abdomen_tip', + 'right_forewing_hinge', 'left_forewing_hinge', 'right_hindwing_hinge', 'left_hindwing_hinge'] + + body_skeleton3d = dict() + for label in labels: + body_skeleton3d[label] = np.asarray(init_dict['init_skeleton3d']['body'][label]) + + body_skeleton3d = build_skeleton3d(body_skeleton3d, body_params) + + return body_skeleton3d + + +def build_skeleton3d(init_body_skeleton3d: Dict[str, List[float]], body_params: Dict[str, float]) -> Dict[str, List[float]]: + """ + + Args: + init_body_skeleton3d: + body_params: + + Returns: + body_skeleton3d: + """ + + body_skeleton3d = copy.deepcopy(init_body_skeleton3d) + + body_params2 = copy.deepcopy(body_params) + for label in ['body_l', 'body_r']: + body_params2[label] = body_params2[label] * body_params2['ratio_body'] + + # Make sure origin is at the center of mass + origin = body_params2['ratio_com_body'] * (body_skeleton3d['head_tip'] - body_skeleton3d['abdomen_tip']) \ + + body_skeleton3d['abdomen_tip'] + + for label in body_skeleton3d.keys(): + body_skeleton3d[label] = body_skeleton3d[label] - origin + + # Set hinges next to the origin (center of mass) + body_skeleton3d['right_forewing_hinge'][0], body_skeleton3d['left_forewing_hinge'][0] = 0.0, 0.0 + body_skeleton3d['right_hindwing_hinge'][0], body_skeleton3d['left_hindwing_hinge'][0] = -0.1, -0.1 + + # Scaling + body_l = body_skeleton3d['head_tip'][0] - body_skeleton3d['abdomen_tip'][0] + body_r = (np.sqrt(body_skeleton3d['right_forewing_hinge'][1] ** 2 + body_skeleton3d['left_forewing_hinge'][1] ** 2) + + np.sqrt(body_skeleton3d['right_hindwing_hinge'][1] ** 2 + body_skeleton3d['left_hindwing_hinge'][1] ** 2))/2 + for label in ['head_tip', 'thorax_middle', 'abdomen_tip', + 'right_forewing_hinge', 'left_forewing_hinge', 'right_hindwing_hinge', 'left_hindwing_hinge']: + body_skeleton3d[label] = np.array([body_skeleton3d[label][0] * body_params2['body_l'] / body_l, + body_skeleton3d[label][1] * body_params2['body_r'] / body_r, + body_skeleton3d[label][2] * body_params2['body_r'] / body_r]) + + body_skeleton3d['thorax_middle'][0] = body_skeleton3d['abdomen_tip'][0] - body_params2['body_l'] / 2 + + return body_skeleton3d + + +def rotate_skeleton3d(body_skeleton3d: Dict[str, List[float]], body_params: Dict[str, float]) -> Dict[str, List[float]]: + """ + + Args: + body_skeleton3d: + body_params: dict with the initial values of the wing parameters (e.g. roll) + + Returns: + body_skeleton3d_world: + """ + # check_body_params_ranges(body_params) + + # Rotation (extrinsic rotations ('xyz') = Rz(yaw_a) @ Ry(pitch_a) @ Rx(roll_a)) + # This is the same as using intrinsic rotation ('ZYX') + # TODO use quaternions? + r_stroke2body = Rotation.from_euler('y', - body_params['default_pitch_a'], degrees=True) + r_world2stroke = Rotation.from_euler('xyz', [- body_params['roll_a'], - body_params['pitch_a'], - body_params['yaw_a']], degrees=True) + + body_skeleton3d_world = copy.deepcopy(body_skeleton3d) + for label in body_skeleton3d_world.keys(): + body_skeleton3d_world[label] = r_world2stroke.apply(r_stroke2body.apply(body_skeleton3d_world[label])) + + return body_skeleton3d_world + + +def translate_skeleton3d(body_skeleton3d: Dict[str, List[float]], body_params: Dict[str, float]) -> Dict[str, List[float]]: + """ + + Args: + body_skeleton3d: + body_params: dict with the initial values of the wing parameters (e.g. roll) + + Returns: + body_skeleton3d_world: + """ + # check_body_params_ranges(body_params) + + # Translating + center_of_mass_coords_world = np.array([body_params['x_com'], body_params['y_com'], body_params['z_com']]) + + body_skeleton3d_world = copy.deepcopy(body_skeleton3d) + for label in body_skeleton3d_world.keys(): + body_skeleton3d_world[label] = body_skeleton3d_world[label] + center_of_mass_coords_world + + return body_skeleton3d_world + + +def estimate_params_from_skeleton3d(body_skeleton3d_world: Dict[str, List[float]], + init_body_params: Dict[str, float]) -> Dict[str, float]: + """ + Estimate parameters of the transformations/rotations needed to get body_skeleton3d_world + + Args: + body_skeleton3d_world: dict with the 3d coordinates of the body skeleton + init_body_params: dict with the initial values of the body parameters (initial transformation/rotation) + + Returns: + body_params: dict with estimated body parameters + """ + + body_params = copy.deepcopy(init_body_params) + body_skeleton3d_world2 = copy.deepcopy(body_skeleton3d_world) + + abdomen_head_vec = body_skeleton3d_world2['abdomen_tip'] - body_skeleton3d_world2['head_tip'] + abdomen_right_hinge_vec = (body_skeleton3d_world2['abdomen_tip'] - + (body_skeleton3d_world2['right_forewing_hinge'] + body_skeleton3d_world2['right_hindwing_hinge'])/2) + abdomen_left_hinge_vec = (body_skeleton3d_world2['abdomen_tip'] - + (body_skeleton3d_world2['left_forewing_hinge'] + body_skeleton3d_world2['left_hindwing_hinge'])/2) + abdomen_com_vec = \ + ((np.dot(abdomen_right_hinge_vec, abdomen_head_vec)/np.sqrt(sum(abdomen_head_vec ** 2)) ** 2) * abdomen_head_vec + + (np.dot(abdomen_left_hinge_vec, abdomen_head_vec) / np.sqrt(sum(abdomen_head_vec ** 2)) ** 2) * abdomen_head_vec)/2 + + body_params['ratio_com_body'] = np.sqrt(sum(abdomen_com_vec ** 2)) / np.sqrt(sum(abdomen_head_vec ** 2)) + + [body_params['x_com'], body_params['y_com'], body_params['z_com']] = \ + body_params['ratio_com_body'] * (body_skeleton3d_world2['head_tip'] - body_skeleton3d_world2['abdomen_tip']) \ + + body_skeleton3d_world2['abdomen_tip'] + + for label in body_skeleton3d_world2.keys(): + body_skeleton3d_world2[label] = \ + body_skeleton3d_world2[label] - [body_params['x_com'], body_params['y_com'], body_params['z_com']] + + # Vectors along the proboscis, torso, abdomen and hinges + body_v = body_skeleton3d_world2['head_tip'] - body_skeleton3d_world2['abdomen_tip'] + hinges_v = ((body_skeleton3d_world2['left_forewing_hinge'] + body_skeleton3d_world2['left_hindwing_hinge'])/2 + - (body_skeleton3d_world2['right_forewing_hinge']) + body_skeleton3d_world2['right_hindwing_hinge'])/2 + + body_v_projxy, body_v_projxz = body_v.copy(), body_v.copy() + body_v_projxy[2], body_v_projxz[1] = 0.0, 0.0 + + # Angles (extrinsic rotations ('xyz') = Rz(yaw_a) @ Ry(pitch_a) @ Rx(roll_a)) + # This is the same as using intrinsic rotation ('ZYX') + body_params['yaw_a'] = - np.sign(body_v_projxy[1]) * get_rotation_angle_btw_vectors(body_v_projxy, [1., 0., 0.]) + + r_stroke2body = Rotation.from_euler('y', - body_params['default_pitch_a'], degrees=True) + r_world2yaw = Rotation.from_euler('z', - body_params['yaw_a'], degrees=True) + + for label in ['pitch_a', 'roll_a']: body_params[label] = np.nan + + if not any(np.isnan(body_v)): + # TODO: Replace .inv by transpose? (orthogonal matrix) + body_v_inv_default_pitch = r_stroke2body.inv().apply(r_world2yaw.inv().apply(body_v)) + body_params['pitch_a'] = np.sign(body_v_inv_default_pitch[2]) * get_rotation_angle_btw_vectors(body_v_inv_default_pitch, [1., 0., 0.]) + + if not any(np.isnan(hinges_v)): + r_world2yaw_pitch = Rotation.from_euler('yz', [- body_params['pitch_a'], - body_params['yaw_a']], degrees=True) + hinges_v_inv_yaw_pitch = r_stroke2body.inv().apply(r_world2yaw_pitch.inv().apply(hinges_v)) + body_params['roll_a'] = - np.sign(hinges_v_inv_yaw_pitch[2]) * get_rotation_angle_btw_vectors(hinges_v_inv_yaw_pitch, [0., 1., 0.]) + + r_world2stroke = Rotation.from_euler('xyz', [- body_params['roll_a'], - body_params['pitch_a'], - body_params['yaw_a']], degrees=True) + + # Round parameters + round_dec = 4 + for label in ['yaw_a', 'pitch_a', 'roll_a']: + body_params[label] = np.round(body_params[label], round_dec) + + # Lengths + body_params['body_l'] = np.round(length_from_vector(body_v), 9) + body_params['body_r'] = np.round(length_from_vector(hinges_v) * np.sqrt(2) / 2, 9) + + # check_body_params_ranges(body_params) + + return body_params + + +def get_copy_params_zero(body_params: Dict[str, float], label_to_zero=['yaw_a', 'pitch_a', 'roll_a']) -> Dict[str, float]: + """ + + Args: + body_params: + label_to_zero: + + Returns: + new_body_params: + """ + new_body_params = copy.deepcopy(body_params) + for label in new_body_params.keys(): + if label in label_to_zero: + new_body_params[label] = 0.0 + + return new_body_params + + +def residuals2d(body_skeleton3d: Dict[str, List[float]], body_skeleton2d: Dict[str, List[float]], + default_residual: float, dlt_coefs: List[List[float]]) -> List[float]: + """ + Compute residuals when comparing body_skeleton3d and body_skeleton2d + + Args: + body_skeleton3d: dict with 3d coordinates of the body skeleton ({'head': [0, 0, 0]}) + body_skeleton2d: dict with 2d coordinates of the body skeleton ({'head': [0, 0, 0]}) + default_residual: float + dlt_coefs: + + Returns: + residuals: distances between body_skeleton3d and body_skeleton2d + """ + nb_cam = len(dlt_coefs) + + body_skeleton2d_from3d = reproject_skeleton3d_to2d(body_skeleton3d, dlt_coefs) + + residuals = np.zeros((len(body_skeleton3d.keys()), nb_cam)) + for i, label in enumerate(body_skeleton3d.keys()): + for j, camn in enumerate(range(1, nb_cam + 1)): + + residuals[i][j] = np.sqrt(np.sum((body_skeleton2d[label][camn][0:2] - body_skeleton2d_from3d[label][camn][0:2]) ** 2)) + + # weight on the likelihood? body_skeleton2d[label][camn][-1] + if np.isnan(residuals[i][j]): + residuals[i][j] = default_residual + + return residuals.flatten() + + +def residuals3d(body_skeleton3d_1: Dict[str, List[float]], body_skeleton3d_2: Dict[str, List[float]], + default_residual: float) -> List[float]: + """ + Compute residuals when comparing body_skeleton3d_1 and body_skeleton3d_2 + + Args: + body_skeleton3d_1: dict with 3d coordinates of the body skeleton ({'head': [0, 0, 0]}) + body_skeleton3d_2: dict with 3d coordinates of the body skeleton ({'head': [0, 0, 0]}) + default_residual: float + + Returns: + residuals: distances between body_skeleton3d_1 and body_skeleton3d_2 + """ + + residuals = np.zeros(len(body_skeleton3d_1.keys())) + for i, label in enumerate(body_skeleton3d_1.keys()): + residuals[i] = np.sqrt(np.sum((body_skeleton3d_2[label] - body_skeleton3d_1[label]) ** 2)) + + if np.isnan(residuals[i]): + residuals[i] = default_residual + + return residuals + + +@dataclass +class BodyModuleSettings: + """ + Class for keeping track of the settings of the insect_body_3segments_2wings module. + """ + + def __init__(self, init_yaml_path: str) -> None: + """ + + Args: + init_yaml_path: Path of yaml file that contains initial kinematic parameters, bounds for these parameters + as well as 3d skeleton coordinates for a given animal (e.g. mosquito) + """ + + with open(init_yaml_path, 'r') as file: + init_dict = yaml.safe_load(file) + + self.params_init = init_params(init_dict) # type: Dict[str, Dict[str, any]] + self.bounds_init = init_bounds(init_dict) # type: Dict[str, Dict[str, any]] + self.skeleton3d_init = init_skeleton3d(init_dict, self.params_init) # type: Dict[str, Dict[str, any]] + + self.param_names = list(self.params_init.keys()) # type: List[str] + self.labels = list(self.skeleton3d_init.keys()) # type: List[str] + + self.param_names_to_keep_cst = ['default_pitch_a', 'ratio_com_body', 'ratio_body', 'body_l', 'body_r'] # type: List[str] + + self.threshold_likelihood = 0.9 # type: float + self.threshold_nb_pts = 4 # type: int + + assert 0.0 < self.threshold_likelihood <= 1.0 + assert 1 < self.threshold_nb_pts <= len(self.labels) + + @staticmethod + def from_dict(init_yaml_path: str, dictionary: Dict[str, any]): + """ + + Args: + init_yaml_path: Path of yaml file that contains initial kinematic parameters, bounds for these parameters + as well as 3d skeleton coordinates for a given animal (e.g. mosquito) + dictionary: Dict that contains values of class attributes that will be updated + """ + + module_sets = BodyModuleSettings(init_yaml_path) + for key, value in dictionary.items(): + setattr(module_sets, key, value) + + return module_sets + diff --git a/skeleton_fitter/modules/bodies/insect_body_slim.py b/skeleton_fitter/modules/bodies/insect_body_3segments_2wings.py similarity index 99% rename from skeleton_fitter/modules/bodies/insect_body_slim.py rename to skeleton_fitter/modules/bodies/insect_body_3segments_2wings.py index 50dc2d3e3285879b7e0a0a3b5e28a89c78c55798..94c3a03dad0069788d427a8f5fc3534f831821f7 100644 --- a/skeleton_fitter/modules/bodies/insect_body_slim.py +++ b/skeleton_fitter/modules/bodies/insect_body_3segments_2wings.py @@ -357,7 +357,7 @@ def residuals3d(body_skeleton3d_1: Dict[str, List[float]], body_skeleton3d_2: Di @dataclass class BodyModuleSettings: """ - Class for keeping track of the settings of the insect_body_slim module. + Class for keeping track of the settings of the insect_body_3segments_2wings module. """ def __init__(self, init_yaml_path: str) -> None: diff --git a/skeleton_fitter/modules/bodies/insect_hybrid_body_slim.py b/skeleton_fitter/modules/bodies/insect_body_3segments_2wings_hybrid.py similarity index 95% rename from skeleton_fitter/modules/bodies/insect_hybrid_body_slim.py rename to skeleton_fitter/modules/bodies/insect_body_3segments_2wings_hybrid.py index 202196d4e8b35a923f5d319eb4243981c3e0b164..71e402dd77050a704be9c36c97f97d10681d82d5 100644 --- a/skeleton_fitter/modules/bodies/insect_hybrid_body_slim.py +++ b/skeleton_fitter/modules/bodies/insect_body_3segments_2wings_hybrid.py @@ -8,10 +8,10 @@ from dataclasses import dataclass from typing import Dict, List from skeleton_fitter.utils import set_axes3d_equal -import skeleton_fitter.modules.bodies.insect_body_slim as body +import skeleton_fitter.modules.bodies.insect_body_3segments_4wings as body # Need these imports to be able to use these functions directly from these hybrid module (without duplicating them here) -from skeleton_fitter.modules.bodies.insect_body_slim import rotate_skeleton3d, translate_skeleton3d, \ +from skeleton_fitter.modules.bodies.insect_body_3segments_4wings import rotate_skeleton3d, translate_skeleton3d, \ estimate_params_from_skeleton3d, get_copy_params_zero, residuals2d, residuals3d @@ -135,7 +135,7 @@ def plot(body_skeleton3d: Dict[str, List[float]], wings_skeleton3d: Dict[str, Li @dataclass class BodyModuleSettings: """ - Class for keeping track of the settings of the insect_hybrid_body_slim module. + Class for keeping track of the settings of the insect_body_3segments_2wings_hybrid module. """ diff --git a/skeleton_fitter/modules/bodies/insect_body_3segments_4wings.py b/skeleton_fitter/modules/bodies/insect_body_3segments_4wings.py new file mode 100644 index 0000000000000000000000000000000000000000..8c4a8e1f0fb3fd7d6cfc7a8137ed300e43b8290a --- /dev/null +++ b/skeleton_fitter/modules/bodies/insect_body_3segments_4wings.py @@ -0,0 +1,413 @@ +import copy, yaml +import numpy as np + +from dataclasses import dataclass + +from scipy.spatial.transform import Rotation + +from typing import Dict, List + +from skeleton_fitter.utils import reproject_skeleton3d_to2d, length_from_vector +from camera.utils import get_rotation_angle_btw_vectors + + +def init_params(init_dict: Dict[str, any]) -> Dict[str, float]: + """ + Load initial kinematic parameters + + Args: + init_dict: Dictionary with initial values of the skeleton coordinates/parameters/bounds + + Returns: + body_params: + """ + + param_names = ['default_pitch_a', 'yaw_a', 'pitch_a', 'roll_a', + 'ratio_com_torso', 'ratio_proboscis_head', 'ratio_body', + 'proboscis_torso_a', 'torso_abdomen_a', 'proboscis_l', 'torso_l', 'torso_r', 'abdomen_l', + 'x_com', 'y_com', 'z_com'] + + body_params = dict() + for param_name in param_names: + body_params[param_name] = init_dict['init_params']['body'][param_name] + + return body_params + + +def init_bounds(init_dict: Dict[str, any]) -> Dict[str, float]: + """ + Load initial bounds for kinematic parameters + + Args: + init_dict: Dictionary with initial values of the skeleton coordinates/parameters/bounds + + Returns: + body_bounds: + """ + + param_names = ['default_pitch_a', 'yaw_a', 'pitch_a', 'roll_a', + 'ratio_com_torso', 'ratio_proboscis_head', 'ratio_body', + 'proboscis_torso_a', 'torso_abdomen_a', 'proboscis_l', 'torso_l', 'torso_r', 'abdomen_l', + 'x_com', 'y_com', 'z_com'] + + body_bounds = dict() + for param_name in param_names: + body_bounds[param_name] = tuple(init_dict['init_bounds']['body'][param_name]) + + return body_bounds + + +def init_skeleton3d(init_dict: Dict[str, any], body_params: Dict[str, float]) -> Dict[str, List[float]]: + """ + Initialize body skeleton + + Args: + init_dict: Dictionary with initial values of the skeleton coordinates/parameters/bounds + body_params: + + Returns: + body_skeleton3d: + """ + + # Initialise a body with 3 segments (proboscis-head, torso, abdomen) along x axis + labels = ['proboscis_tip', 'proboscis_head_joint', 'head_torso_joint', + 'right_forewing_hinge', 'left_forewing_hinge', 'right_hindwing_hinge', 'left_hindwing_hinge', + 'torso_abdomen_joint', 'abdomen_middle', 'abdomen_tip'] + + body_skeleton3d = dict() + for label in labels: + body_skeleton3d[label] = np.asarray(init_dict['init_skeleton3d']['body'][label]) + + body_skeleton3d = build_skeleton3d(body_skeleton3d, body_params) + + return body_skeleton3d + + +def build_skeleton3d(init_body_skeleton3d: Dict[str, List[float]], body_params: Dict[str, float]) -> Dict[str, List[float]]: + """ + + Args: + init_body_skeleton3d: + body_params: + + Returns: + body_skeleton3d: + """ + + body_skeleton3d = copy.deepcopy(init_body_skeleton3d) + + body_params2 = copy.deepcopy(body_params) + for label in ['torso_r', 'torso_l', 'abdomen_l', 'proboscis_l']: + body_params2[label] = body_params2[label] * body_params2['ratio_body'] + + # Make sure origin is at the center of mass + origin = body_params2['ratio_com_torso'] * (body_skeleton3d['head_torso_joint'] - body_skeleton3d['torso_abdomen_joint']) \ + + body_skeleton3d['torso_abdomen_joint'] + + for label in body_skeleton3d.keys(): + body_skeleton3d[label] = body_skeleton3d[label] - origin + + # Set hinges next to the origin (center of mass) + body_skeleton3d['right_forewing_hinge'][0], body_skeleton3d['left_forewing_hinge'][0] = 0.0, 0.0 + body_skeleton3d['right_hindwing_hinge'][0], body_skeleton3d['left_hindwing_hinge'][0] = -0.1, -0.1 + + # Scaling + torso_l = body_skeleton3d['head_torso_joint'][0] - body_skeleton3d['torso_abdomen_joint'][0] + torso_r = (np.sqrt(body_skeleton3d['right_forewing_hinge'][1] ** 2 + body_skeleton3d['left_forewing_hinge'][1] ** 2) + + np.sqrt(body_skeleton3d['right_hindwing_hinge'][1] ** 2 + body_skeleton3d['left_hindwing_hinge'][1] ** 2))/2 + for label in ['head_torso_joint', 'torso_abdomen_joint', + 'right_forewing_hinge', 'left_forewing_hinge', 'right_hindwing_hinge', 'left_hindwing_hinge']: + body_skeleton3d[label] = np.array([body_skeleton3d[label][0] * body_params2['torso_l'] / torso_l, + body_skeleton3d[label][1] * body_params2['torso_r'] / torso_r, + body_skeleton3d[label][2] * body_params2['torso_r'] / torso_r]) + + body_skeleton3d['proboscis_tip'][0] = body_skeleton3d['head_torso_joint'][0] + body_params2['proboscis_l'] + body_skeleton3d['proboscis_head_joint'][0] = \ + body_skeleton3d['head_torso_joint'][0] + body_params2['proboscis_l'] * body_params2['ratio_proboscis_head'] + body_skeleton3d['abdomen_middle'][0] = body_skeleton3d['torso_abdomen_joint'][0] - body_params2['abdomen_l'] / 2 + body_skeleton3d['abdomen_tip'][0] = body_skeleton3d['torso_abdomen_joint'][0] - body_params2['abdomen_l'] + + # Curving + body_skeleton3d['proboscis_tip'][0] = body_skeleton3d['head_torso_joint'][0] \ + + body_params2['proboscis_l'] * np.cos(np.deg2rad(180.0 - body_params2['proboscis_torso_a'])) + body_skeleton3d['proboscis_head_joint'][0] = body_skeleton3d['head_torso_joint'][0] \ + + body_params2['ratio_proboscis_head'] * body_params2['proboscis_l'] * np.cos(np.deg2rad(180.0 - body_params2['proboscis_torso_a'])) + body_skeleton3d['abdomen_middle'][0] = body_skeleton3d['torso_abdomen_joint'][0] \ + - body_params2['abdomen_l'] / 2 * np.cos(np.deg2rad(180.0 - body_params2['torso_abdomen_a'])) + body_skeleton3d['abdomen_tip'][0] = body_skeleton3d['torso_abdomen_joint'][0] \ + - body_params2['abdomen_l'] * np.cos(np.deg2rad(180.0 - body_params2['torso_abdomen_a'])) + + body_skeleton3d['proboscis_tip'][2] = - body_params2['proboscis_l'] * np.sin(np.deg2rad(180.0 - body_params2['proboscis_torso_a'])) + body_skeleton3d['proboscis_head_joint'][2] = \ + - body_params2['ratio_proboscis_head'] * body_params2['proboscis_l'] * np.sin(np.deg2rad(180.0 - body_params2['proboscis_torso_a'])) + body_skeleton3d['abdomen_middle'][2] = - body_params2['abdomen_l'] / 2 * np.sin(np.deg2rad(180.0 - body_params2['torso_abdomen_a'])) + body_skeleton3d['abdomen_tip'][2] = - body_params2['abdomen_l'] * np.sin(np.deg2rad(180.0 - body_params2['torso_abdomen_a'])) + + return body_skeleton3d + + +def rotate_skeleton3d(body_skeleton3d: Dict[str, List[float]], body_params: Dict[str, float]) -> Dict[str, List[float]]: + """ + + Args: + body_skeleton3d: + body_params: dict with the initial values of the wing parameters (e.g. roll) + + Returns: + body_skeleton3d_world: + """ + # check_body_params_ranges(body_params) + + # Rotation (extrinsic rotations ('xyz') = Rz(yaw_a) @ Ry(pitch_a) @ Rx(roll_a)) + # This is the same as using intrinsic rotation ('ZYX') + # TODO use quaternions? + r_stroke2body = Rotation.from_euler('y', - body_params['default_pitch_a'], degrees=True) + r_world2stroke = Rotation.from_euler('xyz', [- body_params['roll_a'], - body_params['pitch_a'], - body_params['yaw_a']], degrees=True) + + body_skeleton3d_world = copy.deepcopy(body_skeleton3d) + for label in body_skeleton3d_world.keys(): + body_skeleton3d_world[label] = r_world2stroke.apply(r_stroke2body.apply(body_skeleton3d_world[label])) + + return body_skeleton3d_world + + +def translate_skeleton3d(body_skeleton3d: Dict[str, List[float]], body_params: Dict[str, float]) -> Dict[str, List[float]]: + """ + + Args: + body_skeleton3d: + body_params: dict with the initial values of the wing parameters (e.g. roll) + + Returns: + body_skeleton3d_world: + """ + # check_body_params_ranges(body_params) + + # Translating + center_of_mass_coords_world = np.array([body_params['x_com'], body_params['y_com'], body_params['z_com']]) + + body_skeleton3d_world = copy.deepcopy(body_skeleton3d) + for label in body_skeleton3d_world.keys(): + body_skeleton3d_world[label] = body_skeleton3d_world[label] + center_of_mass_coords_world + + return body_skeleton3d_world + + +def estimate_params_from_skeleton3d(body_skeleton3d_world: Dict[str, List[float]], + init_body_params: Dict[str, float]) -> Dict[str, float]: + """ + Estimate parameters of the transformations/rotations needed to get body_skeleton3d_world + + Args: + body_skeleton3d_world: dict with the 3d coordinates of the body skeleton + init_body_params: dict with the initial values of the body parameters (initial transformation/rotation) + + Returns: + body_params: dict with estimated body parameters + """ + # Only works if all angles < 90.0 and > - 90.0° (except body_params['proboscis_torso_a'] and body_params['torso_abdomen_a']) + + body_params = copy.deepcopy(init_body_params) + body_skeleton3d_world2 = copy.deepcopy(body_skeleton3d_world) + + head_head_vec = body_skeleton3d_world2['head_torso_joint'] - body_skeleton3d_world2['proboscis_head_joint'] + head_proboscis_vec = body_skeleton3d_world2['head_torso_joint'] - body_skeleton3d_world2['proboscis_tip'] + body_params['ratio_proboscis_head'] = np.sqrt(sum(head_head_vec ** 2)) / np.sqrt(sum(head_proboscis_vec ** 2)) + + abdomen_head_vec = body_skeleton3d_world2['torso_abdomen_joint'] - body_skeleton3d_world2['head_torso_joint'] + abdomen_right_hinge_vec = (body_skeleton3d_world2['torso_abdomen_joint'] - + (body_skeleton3d_world2['right_forewing_hinge'] + body_skeleton3d_world2['right_hindwing_hinge'])/2) + abdomen_left_hinge_vec = (body_skeleton3d_world2['torso_abdomen_joint'] - + (body_skeleton3d_world2['left_forewing_hinge'] + body_skeleton3d_world2['left_hindwing_hinge'])/2) + abdomen_com_vec = \ + ((np.dot(abdomen_right_hinge_vec, abdomen_head_vec)/np.sqrt(sum(abdomen_head_vec ** 2)) ** 2) * abdomen_head_vec + + (np.dot(abdomen_left_hinge_vec, abdomen_head_vec) / np.sqrt(sum(abdomen_head_vec ** 2)) ** 2) * abdomen_head_vec)/2 + + body_params['ratio_com_torso'] = np.sqrt(sum(abdomen_com_vec ** 2)) / np.sqrt(sum(abdomen_head_vec ** 2)) + + [body_params['x_com'], body_params['y_com'], body_params['z_com']] = \ + body_params['ratio_com_torso'] * (body_skeleton3d_world2['head_torso_joint'] + - body_skeleton3d_world2['torso_abdomen_joint']) \ + + body_skeleton3d_world2['torso_abdomen_joint'] + + for label in body_skeleton3d_world2.keys(): + body_skeleton3d_world2[label] = \ + body_skeleton3d_world2[label] - [body_params['x_com'], body_params['y_com'], body_params['z_com']] + + # Vectors along the proboscis, torso, abdomen and hinges + proboscis_v = body_skeleton3d_world2['proboscis_tip'] - body_skeleton3d_world2['head_torso_joint'] + torso_v = body_skeleton3d_world2['head_torso_joint'] - body_skeleton3d_world2['torso_abdomen_joint'] + abdomen_v = body_skeleton3d_world2['torso_abdomen_joint'] - body_skeleton3d_world2['abdomen_tip'] + hinges_v = ((body_skeleton3d_world2['left_forewing_hinge'] + body_skeleton3d_world2['left_hindwing_hinge'])/2 + - (body_skeleton3d_world2['right_forewing_hinge']) + body_skeleton3d_world2['right_hindwing_hinge'])/2 + + torso_v_projxy, torso_v_projxz = torso_v.copy(), torso_v.copy() + torso_v_projxy[2], torso_v_projxz[1] = 0.0, 0.0 + + # Angles (extrinsic rotations ('xyz') = Rz(yaw_a) @ Ry(pitch_a) @ Rx(roll_a)) + # This is the same as using intrinsic rotation ('ZYX') + body_params['yaw_a'] = - np.sign(torso_v_projxy[1]) * get_rotation_angle_btw_vectors(torso_v_projxy, [1., 0., 0.]) + + r_stroke2body = Rotation.from_euler('y', - body_params['default_pitch_a'], degrees=True) + r_world2yaw = Rotation.from_euler('z', - body_params['yaw_a'], degrees=True) + + for label in ['pitch_a', 'roll_a', 'proboscis_torso_a', 'torso_abdomen_a']: body_params[label] = np.nan + + if not any(np.isnan(torso_v)): + # TODO: Replace .inv by transpose? (orthogonal matrix) + torso_v_inv_default_pitch = r_stroke2body.inv().apply(r_world2yaw.inv().apply(torso_v)) + body_params['pitch_a'] = np.sign(torso_v_inv_default_pitch[2]) * get_rotation_angle_btw_vectors(torso_v_inv_default_pitch, [1., 0., 0.]) + + if not any(np.isnan(hinges_v)): + r_world2yaw_pitch = Rotation.from_euler('yz', [- body_params['pitch_a'], - body_params['yaw_a']], degrees=True) + hinges_v_inv_yaw_pitch = r_stroke2body.inv().apply(r_world2yaw_pitch.inv().apply(hinges_v)) + body_params['roll_a'] = - np.sign(hinges_v_inv_yaw_pitch[2]) * get_rotation_angle_btw_vectors(hinges_v_inv_yaw_pitch, [0., 1., 0.]) + + r_world2stroke = Rotation.from_euler('xyz', [- body_params['roll_a'], - body_params['pitch_a'], - body_params['yaw_a']], degrees=True) + + body_params['proboscis_torso_a'] = \ + 180.0 - get_rotation_angle_btw_vectors(r_stroke2body.inv().apply(r_world2stroke.inv().apply(proboscis_v)), + r_stroke2body.inv().apply(r_world2stroke.inv().apply(torso_v))) + body_params['torso_abdomen_a'] = \ + 180.0 - get_rotation_angle_btw_vectors(r_stroke2body.inv().apply(r_world2stroke.inv().apply(torso_v)), + r_stroke2body.inv().apply(r_world2stroke.inv().apply(abdomen_v))) + + # Round parameters + round_dec = 4 + for label in ['yaw_a', 'pitch_a', 'roll_a', 'proboscis_torso_a', 'torso_abdomen_a']: + body_params[label] = np.round(body_params[label], round_dec) + + # Lengths + body_params['proboscis_l'] = np.round(length_from_vector(proboscis_v), 9) + body_params['torso_l'] = np.round(length_from_vector(torso_v), 9) + body_params['torso_r'] = np.round(length_from_vector(hinges_v) * np.sqrt(2) / 2, 9) + body_params['abdomen_l'] = np.round(length_from_vector(abdomen_v), 9) + + # check_body_params_ranges(body_params) + + return body_params + + +def get_copy_params_zero(body_params: Dict[str, float], label_to_zero=['yaw_a', 'pitch_a', 'roll_a']) -> Dict[str, float]: + """ + + Args: + body_params: + label_to_zero: + + Returns: + new_body_params: + """ + new_body_params = copy.deepcopy(body_params) + for label in new_body_params.keys(): + if label in label_to_zero: + new_body_params[label] = 0.0 + + return new_body_params + + +def residuals2d(body_skeleton3d: Dict[str, List[float]], body_skeleton2d: Dict[str, List[float]], + default_residual: float, dlt_coefs: List[List[float]]) -> List[float]: + """ + Compute residuals when comparing body_skeleton3d and body_skeleton2d + + Args: + body_skeleton3d: dict with 3d coordinates of the body skeleton ({'head': [0, 0, 0]}) + body_skeleton2d: dict with 2d coordinates of the body skeleton ({'head': [0, 0, 0]}) + default_residual: float + dlt_coefs: + + Returns: + residuals: distances between body_skeleton3d and body_skeleton2d + """ + nb_cam = len(dlt_coefs) + + body_skeleton2d_from3d = reproject_skeleton3d_to2d(body_skeleton3d, dlt_coefs) + + residuals = np.zeros((len(body_skeleton3d.keys()), nb_cam)) + for i, label in enumerate(body_skeleton3d.keys()): + for j, camn in enumerate(range(1, nb_cam + 1)): + + residuals[i][j] = np.sqrt(np.sum((body_skeleton2d[label][camn][0:2] - body_skeleton2d_from3d[label][camn][0:2]) ** 2)) + + # weight on the likelihood? body_skeleton2d[label][camn][-1] + if np.isnan(residuals[i][j]): + residuals[i][j] = default_residual + + return residuals.flatten() + + +def residuals3d(body_skeleton3d_1: Dict[str, List[float]], body_skeleton3d_2: Dict[str, List[float]], + default_residual: float) -> List[float]: + """ + Compute residuals when comparing body_skeleton3d_1 and body_skeleton3d_2 + + Args: + body_skeleton3d_1: dict with 3d coordinates of the body skeleton ({'head': [0, 0, 0]}) + body_skeleton3d_2: dict with 3d coordinates of the body skeleton ({'head': [0, 0, 0]}) + default_residual: float + + Returns: + residuals: distances between body_skeleton3d_1 and body_skeleton3d_2 + """ + + residuals = np.zeros(len(body_skeleton3d_1.keys())) + for i, label in enumerate(body_skeleton3d_1.keys()): + residuals[i] = np.sqrt(np.sum((body_skeleton3d_2[label] - body_skeleton3d_1[label]) ** 2)) + + if np.isnan(residuals[i]): + residuals[i] = default_residual + + return residuals + + +@dataclass +class BodyModuleSettings: + """ + Class for keeping track of the settings of the insect_body_3segments_2wings module. + """ + + def __init__(self, init_yaml_path: str) -> None: + """ + + Args: + init_yaml_path: Path of yaml file that contains initial kinematic parameters, bounds for these parameters + as well as 3d skeleton coordinates for a given animal (e.g. mosquito) + """ + + with open(init_yaml_path, 'r') as file: + init_dict = yaml.safe_load(file) + + self.params_init = init_params(init_dict) # type: Dict[str, Dict[str, any]] + self.bounds_init = init_bounds(init_dict) # type: Dict[str, Dict[str, any]] + self.skeleton3d_init = init_skeleton3d(init_dict, self.params_init) # type: Dict[str, Dict[str, any]] + + self.param_names = list(self.params_init.keys()) # type: List[str] + self.labels = list(self.skeleton3d_init.keys()) # type: List[str] + + self.param_names_to_keep_cst = ['default_pitch_a', 'ratio_com_torso', 'ratio_proboscis_head', 'ratio_body', + 'proboscis_torso_a', 'torso_abdomen_a', 'proboscis_l', + 'torso_l', 'torso_r', 'abdomen_l'] # type: List[str] + + self.threshold_likelihood = 0.9 # type: float + self.threshold_nb_pts = 4 # type: int + + assert 0.0 < self.threshold_likelihood <= 1.0 + assert 1 < self.threshold_nb_pts <= len(self.labels) + + @staticmethod + def from_dict(init_yaml_path: str, dictionary: Dict[str, any]): + """ + + Args: + init_yaml_path: Path of yaml file that contains initial kinematic parameters, bounds for these parameters + as well as 3d skeleton coordinates for a given animal (e.g. mosquito) + dictionary: Dict that contains values of class attributes that will be updated + """ + + module_sets = BodyModuleSettings(init_yaml_path) + for key, value in dictionary.items(): + setattr(module_sets, key, value) + + return module_sets + diff --git a/skeleton_fitter/modules/limbs/insect_wings_flat.py b/skeleton_fitter/modules/limbs/insect_wings_flat.py index 10a0ec45b0dc43d1aceac3a3afdb9224a750bfa8..551c2bfe4dcf101c41e29c200a678da840629d15 100644 --- a/skeleton_fitter/modules/limbs/insect_wings_flat.py +++ b/skeleton_fitter/modules/limbs/insect_wings_flat.py @@ -8,7 +8,7 @@ from scipy.spatial.transform import Rotation from typing import Dict, List -from skeleton_fitter.modules.bodies.insect_body_slim import BodyModuleSettings +from skeleton_fitter.modules.bodies.insect_body_3segments_4wings import BodyModuleSettings from skeleton_fitter.utils import reproject_skeleton3d_to2d, length_from_vector from camera.utils import get_rotation_angle_btw_vectors diff --git a/skeleton_fitter/modules/limbs/insect_wings_flat_geo.py b/skeleton_fitter/modules/limbs/insect_wings_flat_geo.py index 1dc413750dae22c7b76fe811aee5307c472735cd..7382338cdaa79c9ee2b2e03ba733a2a48d9f3b1a 100644 --- a/skeleton_fitter/modules/limbs/insect_wings_flat_geo.py +++ b/skeleton_fitter/modules/limbs/insect_wings_flat_geo.py @@ -7,7 +7,7 @@ from scipy import stats from typing import Dict, List -from skeleton_fitter.modules.bodies.insect_body_slim import BodyModuleSettings +from skeleton_fitter.modules.bodies.insect_body_3segments_4wings import BodyModuleSettings from skeleton_fitter.utils import reproject_skeleton3d_to2d, length_from_vector diff --git a/skeleton_fitter/modules/limbs/insect_wings_flat_with_apex.py b/skeleton_fitter/modules/limbs/insect_wings_flat_with_apex.py new file mode 100644 index 0000000000000000000000000000000000000000..e5bead761272e0b58e0c65071f3240ffb0c9a5d2 --- /dev/null +++ b/skeleton_fitter/modules/limbs/insect_wings_flat_with_apex.py @@ -0,0 +1,450 @@ +import copy, yaml +import numpy as np + +from dataclasses import dataclass + + +from scipy.spatial.transform import Rotation + +from typing import Dict, List + +from skeleton_fitter.modules.bodies.insect_body_3segments_4wings import BodyModuleSettings + +from skeleton_fitter.utils import reproject_skeleton3d_to2d, length_from_vector +from camera.utils import get_rotation_angle_btw_vectors + + +def init_params(init_dict: Dict[str, any], body_params: Dict[str, float]) -> Dict[str, float]: + """ + + Args: + init_dict: Dictionary with initial values of the skeleton coordinates/parameters/bounds + body_params: + + Returns: + wings_params: + """ + + param_names_to_load_from_yaml = ['stroke_a', 'deviation_a', 'rotation_a', 'span', 'chord', 'ratio_wing'] + + wings_params = {'right': copy.deepcopy(body_params), 'left': copy.deepcopy(body_params)} + for side in ['right', 'left']: + for param_name in param_names_to_load_from_yaml: + wings_params[side][param_name] = init_dict['init_params']['limbs']['wing'][param_name] + + # TODO! Add the possibility to estimate wings' curvatures + # wings_params[side]['span_curvature'] = 0.0 # curvature = 1/arc_length + # wings_params[side]['chord_curvature'] = 0.0 + + wings_params[side]['aspect_ratio'] = wings_params[side]['span'] / wings_params[side]['chord'] + + wings_params[side]['x_hinge'] = body_params['x_com'] + if side == 'right': wings_params[side]['y_hinge'] = body_params['y_com'] - np.sqrt(2) * body_params['torso_r'] + elif side == 'left': wings_params[side]['y_hinge'] = body_params['y_com'] + np.sqrt(2) * body_params['torso_r'] + wings_params[side]['z_hinge'] = body_params['z_com'] + np.sqrt(2) * body_params['torso_r'] + + return wings_params + + +def init_bounds(init_dict: Dict[str, any], body_bounds: Dict[str, float]) -> Dict[str, float]: + """ + + Args: + init_dict: Dictionary with initial values of the skeleton coordinates/parameters/bounds + body_bounds: + + Returns: + wing_bounds: + """ + + param_names = ['stroke_a', 'deviation_a', 'rotation_a', 'span', 'chord', 'ratio_wing', 'aspect_ratio', + 'x_hinge', 'y_hinge', 'z_hinge'] + + wing_bounds = copy.deepcopy(body_bounds) + for param_name in param_names: + wing_bounds[param_name] = tuple(init_dict['init_bounds']['limbs']['wing'][param_name]) + + # TODO! Add the possibility to estimate wings' curvatures + # wing_bounds['span_curvature'] = (0.0, 1.0) + # wing_bounds['chord_curvature'] = (0.0, 1.0) + + return wing_bounds + + +def init_skeleton3d(init_dict: Dict[str, any], wings_params: Dict[str, float]) -> Dict[str, List[float]]: + """ + + Args: + init_dict: Dictionary with initial values of the skeleton coordinates/parameters/bounds + wings_params: + + Returns: + wings_skeleton3d: + """ + + labels = ['hinge', 'leading_edge_q1', 'leading_edge_q2', 'leading_edge_q3', + 'apex', 'external_edge_q1', 'external_edge_q2', 'external_edge_q3', + 'tip', 'trailing_edge_q3', 'trailing_edge_q2', 'trailing_edge_q1'] + + wings_skeleton3d = {} + for side in ['left', 'right']: + + wings_skeleton3d[side] = {} + for label in labels: + wings_skeleton3d[side][label] = np.asarray(init_dict['init_skeleton3d']['limbs']['wing'][label]) + + if side == 'right': # the default wing is the left wing, the right need to be mirrored along th y axis + wings_skeleton3d[side][label][1] = -wings_skeleton3d[side][label][1] + + wings_skeleton3d[side] = build_skeleton3d(wings_skeleton3d[side], wings_params[side]) + + return wings_skeleton3d + + +def build_skeleton3d(init_wing_skeleton3d: Dict[str, List[float]], wing_params: Dict[str, float]) -> Dict[str, List[float]]: + """ + + Args: + init_wing_skeleton3d: + wing_params: + + Returns: + wing_skeleton3d: + """ + wing_skeleton3d = copy.deepcopy(init_wing_skeleton3d) + + # Scaling + x_scaling_factor = wing_params['chord'] / np.abs(wing_skeleton3d['leading_edge_q2'][0] - wing_skeleton3d['trailing_edge_q2'][0]) + y_scaling_factor = wing_params['span'] / np.abs(wing_skeleton3d['tip'][1] - wing_skeleton3d['hinge'][1]) + + for label in wing_skeleton3d.keys(): + wing_skeleton3d[label][0] = x_scaling_factor * wing_skeleton3d[label][0] * wing_params['ratio_wing'] + wing_skeleton3d[label][1] = y_scaling_factor * wing_skeleton3d[label][1] * wing_params['ratio_wing'] + + # # Curving + # ???? + + # Make move wing to hinge (from wing_params) + hinge_coords_world = np.array([wing_params['x_hinge'], wing_params['y_hinge'], wing_params['z_hinge']]) + translation_vector = hinge_coords_world - wing_skeleton3d['hinge'] + for label in wing_skeleton3d.keys(): + wing_skeleton3d[label] = wing_skeleton3d[label] + translation_vector + + return wing_skeleton3d + + +def rotate_and_translate_skeleton3d(wing_skeleton3d: Dict[str, List[float]], wing_params: Dict[str, float], side: str) -> Dict[str, List[float]]: + """ + + Args: + wing_skeleton3d: (in body reference frame) + wing_params: dict with the initial values of the wing parameters (e.g. stroke_a) + side: 'left' or 'right' + + Returns: + wing_skeleton3d_world: + """ + # check_body_params_ranges(body_params) + # check_wings_params_ranges(wings_params) + + # TODO make rotate_and_translate_skeleton3d about only rotating (already the case?)? + + # Rotating and translating + # (extrinsic rotations ('xyz') = Rz(yaw_a) @ Ry(pitch_a) @ Rx(roll_a)) + # This is the same as using intrinsic rotation ('ZYX') + # TODO use quaternions? + r_world2stroke = Rotation.from_euler('xyz', [- wing_params['roll_a'], - wing_params['pitch_a'],- wing_params['yaw_a']], degrees=True) + + if side == 'right': + r_stroke2wing = Rotation.from_euler('yxz', [- wing_params['rotation_a'], - wing_params['deviation_a'], + wing_params['stroke_a']], degrees=True) + elif side == 'left': + r_stroke2wing = Rotation.from_euler('yxz', [- wing_params['rotation_a'], wing_params['deviation_a'], + - wing_params['stroke_a']], degrees=True) + + hinge_coords_init = wing_skeleton3d['hinge'] + hinge_coords_world = np.array([wing_params['x_hinge'], wing_params['y_hinge'], wing_params['z_hinge']]) + center_of_mass_coords_world = np.array([wing_params['x_com'], wing_params['y_com'], wing_params['z_com']]) + + # hinge_coords_stroke = r_world2stroke.inv().apply(r_stroke2body.inv().apply(hinge_coords_world - center_of_mass_coords_world)) + hinge_coords_stroke = r_world2stroke.inv().apply(hinge_coords_world - center_of_mass_coords_world) + + wing_skeleton3d_world = copy.deepcopy(wing_skeleton3d) + for label in wing_skeleton3d_world.keys(): + wing_skeleton3d_world[label] = wing_skeleton3d_world[label] - hinge_coords_init # Set origin to hinge + wing_skeleton3d_world[label] = r_stroke2wing.apply(wing_skeleton3d_world[label]) # Rotate wing into stroke plane + + # Translate and rotate wing to body reference frame (origin at center of mass) + wing_skeleton3d_world[label] = wing_skeleton3d_world[label] + hinge_coords_stroke + + wing_skeleton3d_world[label] = r_world2stroke.apply(wing_skeleton3d_world[label]) # Rotate into world ref. plane + wing_skeleton3d_world[label] = wing_skeleton3d_world[label] + center_of_mass_coords_world # Translate to center of mass in world ref. plane + + return wing_skeleton3d_world + + +def translate_skeleton3d(wing_skeleton3d: Dict[str, List[float]], wing_params: Dict[str, float]) -> Dict[str, List[float]]: + """ + + Args: + wing_skeleton3d: + wing_params: dict with the initial values of the wing parameters (e.g. stroke_a) + hinge_coords_world: + + Returns: + wing_skeleton3d_world: + """ + # check_body_params_ranges(body_params) + + # Translating + hinge_coords_world = np.array([wing_params['x_hinge'], wing_params['y_hinge'], wing_params['z_hinge']]) + + wing_skeleton3d_world = copy.deepcopy(wing_skeleton3d) + for label in wing_skeleton3d_world.keys(): + wing_skeleton3d_world[label] = wing_skeleton3d_world[label] + hinge_coords_world + + return wing_skeleton3d_world + + +def estimate_params_from_skeleton3d(wing_skeleton3d_world: Dict[str, List[float]], body_params: Dict[str, float], + init_wing_params: Dict[str, float], side: str) -> Dict[str, float]: + """ + Estimate parameters of the transformations/rotations needed to get wing_skeleton3d_world + + Args: + wing_skeleton3d_world: dict with the 3d coordinates of the wing skeleton + body_params: previously estimated parameters for the body + init_wing_params: dict with the initial values of the wing parameters (initial transformation/rotation) + side: 'left' or 'right' + + Returns: + wing_params: dict with estimated wing parameters + """ + # Only works if all wings angles < 90.0 and > - 90.0° + + wing_params = copy.deepcopy(init_wing_params) + wing_skeleton_stroke_plan = copy.deepcopy(wing_skeleton3d_world) + for label in body_params.keys(): + wing_params[label] = body_params[label] + + [wing_params['x_hinge'], wing_params['y_hinge'], wing_params['z_hinge']] = wing_skeleton_stroke_plan['hinge'] + for label in wing_skeleton_stroke_plan.keys(): + wing_skeleton_stroke_plan[label] = wing_skeleton_stroke_plan[label] \ + - [wing_params['x_hinge'], wing_params['y_hinge'], wing_params['z_hinge']] + + # Rotation from world ref. frame to stroke ref. frame + # (extrinsic rotations ('xyz') = Rz(yaw_a) @ Ry(pitch_a) @ Rx(roll_a)) + # This is the same as using intrinsic rotation ('ZYX') + # TODO use quaternions? + r_world2stroke = Rotation.from_euler('xyz', [- body_params['roll_a'], - body_params['pitch_a'], - body_params['yaw_a']], degrees=True) + + for label in wing_skeleton_stroke_plan.keys(): + if not any(np.isnan(wing_skeleton_stroke_plan[label])): + wing_skeleton_stroke_plan[label] = r_world2stroke.inv().apply(wing_skeleton_stroke_plan[label]) + + # Vectors along the span and chord + span_v = wing_skeleton_stroke_plan['tip'] - wing_skeleton_stroke_plan['hinge'] + chord_v = wing_skeleton_stroke_plan['leading_edge_q2'] - wing_skeleton_stroke_plan['trailing_edge_q2'] + + span_v_projxy = span_v.copy() + span_v_projxy[2] = 0.0 + + for label in ['stroke_a', 'deviation_a', 'rotation_a']: wing_params[label] = np.nan + + if not any(np.isnan(span_v)) and not any(np.isnan(span_v_projxy)): + # Angles (extrinsic rotations ('xyz') = Rz(yaw_a) @ Ry(pitch_a) @ Rx(roll_a)) + # This is the same as using intrinsic rotation ('ZYX') + # TODO use quaternions? + if side == 'right': + wing_params['stroke_a'] = np.sign(span_v_projxy[0]) * get_rotation_angle_btw_vectors(span_v_projxy, [0., -1., 0.]) + span_v_inv_stroke = Rotation.from_euler('z', wing_params['stroke_a'], degrees=True).inv().apply(span_v) + wing_params['deviation_a'] = np.sign(span_v_inv_stroke[2]) * get_rotation_angle_btw_vectors(span_v_inv_stroke, [0., -1., 0.]) + + # Rotation from world ref. frame to ref. frame before rotation + r_stroke_deviation = Rotation.from_euler('xz', [- wing_params['deviation_a'], wing_params['stroke_a']], degrees=True) + + elif side == 'left': + wing_params['stroke_a'] = np.sign(span_v_projxy[0]) * get_rotation_angle_btw_vectors(span_v_projxy, [0., 1., 0.]) + span_v_inv_stroke = Rotation.from_euler('z', - wing_params['stroke_a'], degrees=True).inv().apply(span_v) + wing_params['deviation_a'] = np.sign(span_v_inv_stroke[2]) * get_rotation_angle_btw_vectors(span_v_inv_stroke, [0., 1., 0.]) + + # Rotation from world ref. frame to ref. frame before rotation + r_stroke_deviation = Rotation.from_euler('xz', [- wing_params['deviation_a'], - wing_params['stroke_a']], degrees=True) + + if not any(np.isnan(chord_v)): + chord_v_inv_stroke_deviation = r_stroke_deviation.inv().apply(chord_v) + wing_params['rotation_a'] = \ + np.sign(chord_v_inv_stroke_deviation[2]) * get_rotation_angle_btw_vectors(chord_v_inv_stroke_deviation, [1., 0., 0.]) + + # Round parameters + round_dec = 4 + for label in ['stroke_a', 'deviation_a', 'rotation_a']: + wing_params[label] = np.round(wing_params[label], round_dec) + + # check_body_params_ranges(body_params) + # check_wings_params_ranges(wings_params) + + # Lengths + wing_params['span'] = np.round(length_from_vector(span_v), 9) + wing_params['chord'] = np.round(length_from_vector(chord_v), 9) + wing_params['aspect_ratio'] = wing_params['span'] / wing_params['chord'] + + # # Curvatures + # wing_params['span_curvature'], wing_params['chord_curvature'] + + return wing_params + + +def residuals2d(wing_skeleton3d: Dict[str, List[float]], wing_skeleton2d: Dict[str, List[float]], + default_residual: float, dlt_coefs: List[List[float]]) -> List[float]: + """ + Compute residuals when comparing wing_skeleton3d to wing_skeleton2d + + Args: + wing_skeleton3d: + wing_skeleton2d: + default_residual: + dlt_coefs: + + Returns: + residuals: distances between wing_skeleton3d and wing_skeleton2d + """ + nb_cam = len(dlt_coefs) + + wing_skeleton2d_from3d = reproject_skeleton3d_to2d(wing_skeleton3d, dlt_coefs) + + residuals = np.zeros((len(wing_skeleton3d.keys()), nb_cam)) + for i, label in enumerate(wing_skeleton3d.keys()): + for j, camn in enumerate(range(1, nb_cam + 1)): + residuals[i][j] = np.sqrt(np.sum((wing_skeleton2d[label][camn][0:2] - wing_skeleton2d_from3d[label][camn][0:2]) ** 2)) + + # weight on the likelihood? wing_skeleton2d[label][camn][-1] + if np.isnan(residuals[i][j]): + residuals[i][j] = default_residual + + return residuals.flatten() + + +def residuals3d(wing_skeleton3d_1: Dict[str, List[float]], wing_skeleton3d_2: Dict[str, List[float]], default_residual: float) -> List[float]: + """ + Compute residuals when comparing wing_skeleton3d_1 to wing_skeleton3d_2 + + Args: + wing_skeleton3d_1: + wing_skeleton3d_2: + default_residual: + + Returns: + residuals: distances between wing_skeleton3d_1 and wing_skeleton3d_2 + """ + residuals = np.zeros(len(wing_skeleton3d_1.keys())) + for i, label in enumerate(wing_skeleton3d_1.keys()): + residuals[i] = np.sqrt(np.sum((wing_skeleton3d_2[label] - wing_skeleton3d_1[label]) ** 2)) + + if np.isnan(residuals[i]): + residuals[i] = default_residual + + return residuals + + +def postprocessing_estimate_limbs_params(wings_params: Dict[str, float], wings_sets): + """ + Animal specific postprocessing function: Here will keep aspect ratio constant if specified in wings_sets + + Args: + wings_params: + wings_sets: + + Returns: + wings_params: + + """ + + if wings_sets.keep_init_aspect_ratio: + for side in ['right', 'left']: + span_v = wings_sets.skeleton3d_init[side]['tip'] - wings_sets.skeleton3d_init[side]['hinge'] + chord_v = wings_sets.skeleton3d_init[side]['leading_edge_q2'] - wings_sets.skeleton3d_init[side]['trailing_edge_q2'] + + init_aspect_ratio = np.round(length_from_vector(span_v), 9) / np.round(length_from_vector(chord_v), 9) + init_aspect_ratio = init_aspect_ratio * np.ones(np.shape(wings_params[side]['span'])) + + wings_params[side]['aspect_ratio'] = init_aspect_ratio + wings_params[side]['chord'] = wings_params[side]['span'] / init_aspect_ratio + + return wings_params + + +def get_copy_multi_params_zero(wings_params: Dict[str, float], + label_to_zero=['stroke_a', 'deviation_a', 'rotation_a']) -> Dict[str, float]: + """ + + Args: + wings_params: + label_to_zero: + + Returns: + new_wings_params: + """ + new_wings_params = copy.deepcopy(wings_params) + for side in ['right', 'left']: + for label in new_wings_params[side].keys(): + if label in label_to_zero: + new_wings_params[side][label] = 0.0 + + return new_wings_params + + +@dataclass +class LimbsModuleSettings: + """ + Class for keeping track of the settings of the insect_wings_flat module. + """ + + def __init__(self, init_yaml_path: str) -> None: + """ + + Args: + init_yaml_path: Path of yaml file that contains initial kinematic parameters, bounds for these parameters + as well as 3d skeleton coordinates for a given animal (e.g. mosquito) + """ + + self.name = 'wing' # type: str + self.sides = ['left', 'right'] # type: List[str] + + self.body_sets = BodyModuleSettings(init_yaml_path) # type: classmethod + + with open(init_yaml_path, 'r') as file: + init_dict = yaml.safe_load(file) + + self.params_init = init_params(init_dict, self.body_sets.params_init) # type: Dict[str, Dict[str, any]] + self.bounds_init = init_bounds(init_dict, self.body_sets.bounds_init) # type: Dict[str, Dict[str, any]] + self.skeleton3d_init = init_skeleton3d(init_dict, self.params_init) # type: Dict[str, Dict[str, any]] + + self.param_names = list(self.params_init[self.sides[0]].keys()) # type: List[str] + self.labels = list(self.skeleton3d_init[self.sides[0]].keys()) # type: List[str] + + self.param_names_to_keep_cst = ['span', 'chord', 'aspect_ratio', 'ratio_wing'] # type: List[str] + + self.threshold_likelihood = 0.9 # type: float + self.threshold_nb_pts = 4 # type: int + + self.keep_init_aspect_ratio = False # type: bool + + assert 0.0 < self.threshold_likelihood <= 1.0 + assert 1 < self.threshold_nb_pts + + @staticmethod + def from_dict(init_yaml_path: str, dictionary: Dict[str, any]): + """ + + Args: + init_yaml_path: Path of yaml file that contains initial kinematic parameters, bounds for these parameters + as well as 3d skeleton coordinates for a given animal (e.g. mosquito) + dictionary: Dict that contains values of class attributes that will be updated + """ + + module_sets = LimbsModuleSettings(init_yaml_path) + for key, value in dictionary.items(): + setattr(module_sets, key, value) + + return module_sets diff --git a/skeleton_fitter/optimiser.py b/skeleton_fitter/optimiser.py index e87879c9eb18d0011a51cc8932cfa351d3b1e12d..51ff86bd4cda7c9f2ca08f58537d305c5554c6ed 100755 --- a/skeleton_fitter/optimiser.py +++ b/skeleton_fitter/optimiser.py @@ -114,7 +114,7 @@ def optim_fit_body_params(module, body_skeleton3d_init, body_skeleton, body_para Args: module: body or limb module that contains the function related to this part of the skeleton - (e.g. import_module(skeleton_fitter.modules.bodies.insect_body_slim)) + (e.g. import_module(skeleton_fitter.modules.bodies.insect_body_3segments_2wings)) body_skeleton3d_init: body_skeleton: body_params_init: