Merge pull request #115 from Skylark0924/czli-dev2

Add human manipulability and ergonomics calculation

examples/visualab/example_ergo_manip.py

@@ -0,0 +1,81 @@
+"""
+Ergo manipulation
+===========================
+
+Examples
+"""
+
+import numpy as np
+import rofunc as rf
+import matplotlib.pyplot as plt
+from matplotlib.animation import FuncAnimation
+
+
+def update(frame):
+    ax.clear()
+    ax.set_xlim(-1, 1)
+    ax.set_ylim(-1, 1)
+    ax.set_zlim(0, 2)
+    ax.set_xlabel('X')
+    ax.set_ylabel('Y')
+    ax.set_zlabel('Z')
+    joint_rotations = skeleton_joint[frame]
+    global_positions, global_rotations = rf.maniplab.forward_kinematics(skeleton_joint_local_translation, joint_rotations, skeleton_parent_indices)
+    rf.visualab.plot_skeleton(ax, global_positions, skeleton_parent_indices)
+    right_hand_index = 5
+    left_hand_index = 8
+    jacobian_right = np.squeeze(jacobians_right[frame])
+    eigenvalues_right, eigenvectors_right = rf.maniplab.calculate_manipulability(jacobian_right)
+    rf.visualab.plot_ellipsoid(ax, eigenvalues_right, eigenvectors_right, global_positions[right_hand_index], 'b')
+    jacobian_left = np.squeeze(jacobians_left[frame])
+    eigenvalues_left, eigenvectors_left = rf.maniplab.calculate_manipulability(jacobian_left)
+    rf.visualab.plot_ellipsoid(ax, eigenvalues_left, eigenvectors_left, global_positions[left_hand_index], 'r')
+
+    rf.visualab.plot_skeleton(ax, global_positions, skeleton_parent_indices)
+
+    # Define a structure to map joint types to their corresponding functions and indices
+    joint_analysis = {
+        'upper_arm': {
+            'indices': [skeleton_joint_name.index('right_upper_arm'), skeleton_joint_name.index('left_upper_arm')],
+            'degree_func': rf.ergolab.UpperArmDegree(global_positions).upper_arm_degrees,
+            'reba_func': rf.ergolab.UAREBA,
+            'score_func': lambda reba: reba.upper_arm_reba_score(),
+            'color_func': rf.visualab.ua_get_color
+        },
+        'lower_arm': {
+            'indices': [skeleton_joint_name.index('right_lower_arm'), skeleton_joint_name.index('left_lower_arm')],
+            'degree_func': rf.ergolab.LADegrees(global_positions).lower_arm_degree,
+            'reba_func': rf.ergolab.LAREBA,
+            'score_func': lambda reba: reba.lower_arm_score(),
+            'color_func': rf.visualab.la_get_color
+        },
+        'trunk': {
+            'indices': [skeleton_joint_name.index('pelvis')],
+            'degree_func': rf.ergolab.TrunkDegree(global_positions, global_rotations).trunk_degrees,
+            'reba_func': rf.ergolab.TrunkREBA,
+            'score_func': lambda reba: reba.trunk_reba_score(),
+            'color_func': rf.visualab.trunk_get_color
+        }
+    }
+
+    # Process each joint type with its corresponding REBA analysis and coloring
+    for joint_type, settings in joint_analysis.items():
+        degrees = settings['degree_func']()
+        reba = settings['reba_func'](degrees)
+        scores = settings['score_func'](reba)
+        for i, idx in enumerate(settings['indices']):
+            score = scores[i]
+            color = settings['color_func'](score)
+            ax.scatter(*global_positions[idx], color=color, s=100)  # Use scatter for single points
+
+
+if __name__ == '__main__':
+    skeleton_joint_name, skeleton_joint, skeleton_parent_indices, skeleton_joint_local_translation = rf.maniplab.read_skeleton_motion('/home/ubuntu/Ergo-Manip/data/demo_2_test_chenzui_only_optitrack2hotu.npy')
+    skeleton_joint = skeleton_joint[::40, :, :]
+    jacobians_left = np.load('/home/ubuntu/Ergo-Manip/data/jacobian_data/jacobians_left_hand_2.npy', allow_pickle=True)[::10]
+    jacobians_right = np.load('/home/ubuntu/Ergo-Manip/data/jacobian_data/jacobians_right_hand_2.npy', allow_pickle=True)[::10]
+    fig = plt.figure()
+    ax = fig.add_subplot(111, projection='3d')
+    ani = FuncAnimation(fig, update, frames=len(skeleton_joint), repeat=True)
+    # ani.save('/home/ubuntu/Ergo-Manip/data/gif/demo_2_andrew.gif', writer='imagemagick', fps=3)
+    plt.show()

rofunc/__init__.py

@@ -22,7 +22,7 @@
 from .learning import ml
 from .learning import RofuncIL, RofuncRL
 from .planning_control import lqt, lqr
-from .utils import visualab, robolab, logger, oslab
+from .utils import visualab, robolab, logger, oslab, ergolab, maniplab
 from .utils.datalab import primitive, data_generator
 from . import config
 

rofunc/utils/ergolab/__init__.py

@@ -0,0 +1,9 @@
+from .reba_calculate.degree2reba.upper_arm_reba_score import *
+from .reba_calculate.degree2reba.trunk_reba_score import *
+from .reba_calculate.degree2reba.lower_arm_reba_score import *
+
+from .reba_calculate.pose2degree.trunk_degree import *
+from .reba_calculate.pose2degree.lower_arm_degree import *
+from .reba_calculate.pose2degree.upper_arm_degree import *
+from .reba_calculate.pose2degree.Util import *
+from .reba_calculate.pose2degree.body_part_sorting import *

rofunc/utils/ergolab/reba_calculate/degree2reba/lower_arm_reba_score.py

@@ -0,0 +1,49 @@
+class LAREBA:
+    # For calculating REBA score based on degrees
+    def __init__(self, LA_degrees):
+        self.LA_degrees = LA_degrees
+
+    def lower_arm_score_overall(self):
+        degree = self.LA_degrees
+        right_degree = degree[0]
+        left_degree = degree[1]
+        lower_arm_reba_score = 0
+        if right_degree >= left_degree:
+            if 0 <= right_degree < 60:
+                lower_arm_reba_score = lower_arm_reba_score + 2
+            if 60 <= right_degree < 100:
+                lower_arm_reba_score = lower_arm_reba_score + 1
+            if 100 <= right_degree:
+                lower_arm_reba_score = lower_arm_reba_score + 1
+        if right_degree < left_degree:
+            if 0 <= left_degree < 60:
+                lower_arm_reba_score = lower_arm_reba_score + 2
+            if 60 <= left_degree < 100:
+                lower_arm_reba_score = lower_arm_reba_score + 1
+            if 100 <= left_degree:
+                lower_arm_reba_score = lower_arm_reba_score + 1
+
+        return lower_arm_reba_score
+
+    def lower_arm_score(self):
+        degree = self.LA_degrees
+        right_degree = degree[0]
+        left_degree = degree[1]
+
+        right_lower_arm_reba_score = 0
+        if 0 <= right_degree < 60:
+            right_lower_arm_reba_score = right_lower_arm_reba_score + 2
+        if 60 <= right_degree < 100:
+            right_lower_arm_reba_score = right_lower_arm_reba_score + 1
+        if 100 <= right_degree:
+            right_lower_arm_reba_score = right_lower_arm_reba_score + 1
+
+        left_lower_arm_reba_score = 0
+        if 0 <= left_degree < 60:
+            left_lower_arm_reba_score = left_lower_arm_reba_score + 2
+        if 60 <= left_degree < 100:
+            left_lower_arm_reba_score = left_lower_arm_reba_score + 1
+        if 100 <= left_degree:
+            left_lower_arm_reba_score = left_lower_arm_reba_score + 1
+
+        return right_lower_arm_reba_score, left_lower_arm_reba_score

rofunc/utils/ergolab/reba_calculate/degree2reba/trunk_reba_score.py

@@ -0,0 +1,50 @@
+class TrunkREBA:
+    # For calculating REBA score based on degrees
+    def __init__(self, trunk_degrees):
+        self.trunk_degrees = trunk_degrees
+
+    def trunk_reba_score(self):
+
+        trunk_flex_degree = self.trunk_degrees[0]
+        trunk_side_bending_degree = self.trunk_degrees[1]
+        trunk_torsion_degree = self.trunk_degrees[2]
+
+        trunk_reba_score = 0
+        trunk_flex_score = 0
+        trunk_side_score = 0
+        trunk_torsion_score = 0
+
+        if trunk_flex_degree >= 0:
+            # means flexion
+            if 0 <= trunk_flex_degree < 5:
+                trunk_reba_score += 1
+                trunk_flex_score += 1
+            elif 5 <= trunk_flex_degree < 20:
+                trunk_reba_score += 2
+                trunk_flex_score += 2
+            elif 20 <= trunk_flex_degree < 60:
+                trunk_reba_score += 3
+                trunk_flex_score += 3
+            elif 60 <= trunk_flex_degree:
+                trunk_reba_score += 4
+                trunk_flex_score += 4
+        else:
+            # means extension
+            if 0 <= abs(trunk_flex_degree) < 5:
+                trunk_reba_score += 1
+                trunk_flex_score += 1
+            elif 5 <= abs(trunk_flex_degree) < 20:
+                trunk_reba_score += 2
+                trunk_flex_score += 2
+            elif 20 <= abs(trunk_flex_degree):
+                trunk_reba_score += 3
+                trunk_flex_score += 3
+
+        if abs(trunk_side_bending_degree) >= 3:
+            trunk_reba_score += 1
+            trunk_side_score += 1
+        if abs(trunk_torsion_degree) >= 1:
+            trunk_reba_score += 1
+            trunk_torsion_score += 1
+
+        return [trunk_reba_score, trunk_flex_score, trunk_side_score, trunk_torsion_score]

rofunc/utils/ergolab/reba_calculate/degree2reba/upper_arm_reba_score.py

@@ -0,0 +1,137 @@
+class UAREBA:
+    # For calculating REBA score based on degrees
+    def __init__(self, UA_degrees):
+        self.UA_degrees = UA_degrees
+
+    def upper_arm_reba_score_overall(self):
+        upper_arm_reba_score = 0
+        upper_arm_flex_score = 0
+        upper_arm_side_score = 0
+        upper_arm_shoulder_rise = 0
+
+
+        right_flexion = self.UA_degrees[0]
+        left_flexion = self.UA_degrees[1]
+
+        right_side = self.UA_degrees[2]
+        left_side = self.UA_degrees[3]
+
+        right_shoulder_rise = self.UA_degrees[4]
+        left_shoulder_rise = self.UA_degrees[5]
+
+        if right_flexion >= left_flexion:
+            if -20 <= right_flexion < 20:
+                upper_arm_reba_score += 1
+                upper_arm_flex_score += 1
+            if 20 <= right_flexion < 45:
+                upper_arm_reba_score += 2
+                upper_arm_flex_score += 2
+            if right_flexion < -20:
+                upper_arm_reba_score += 2
+                upper_arm_flex_score += 2
+            if 45 <= right_flexion < 90:
+                upper_arm_reba_score += 3
+                upper_arm_flex_score += 3
+            if 90 <= right_flexion:
+                upper_arm_reba_score += 4
+                upper_arm_flex_score += 4
+        if right_flexion < left_flexion:
+            if -20 <= left_flexion < 20:
+                upper_arm_reba_score += 1
+                upper_arm_flex_score += 1
+            if left_flexion < -20:
+                upper_arm_reba_score += 2
+                upper_arm_flex_score += 2
+            if 20 <= left_flexion < 45:
+                upper_arm_reba_score += 2
+                upper_arm_flex_score += 2
+            if 45 <= left_flexion < 90:
+                upper_arm_reba_score += 3
+                upper_arm_flex_score += 3
+            if 90 <= left_flexion:
+                upper_arm_reba_score += 4
+                upper_arm_flex_score += 4
+
+            # for side bending
+            if abs(right_side) > 2 or abs(left_side) > 2:
+                upper_arm_reba_score += 1
+                upper_arm_side_score += 1
+
+            # for shoulder rise
+
+            if right_shoulder_rise > 90 or left_shoulder_rise > 90:
+                upper_arm_reba_score += 1
+                upper_arm_shoulder_rise += 1
+        return [upper_arm_reba_score, upper_arm_flex_score, upper_arm_side_score, upper_arm_shoulder_rise]
+
+    def upper_arm_reba_score(self):
+        right_upper_arm_reba_score = 0
+        right_upper_arm_flex_score = 0
+        right_upper_arm_side_score = 0
+        right_upper_arm_shoulder_rise = 0
+        left_upper_arm_reba_score = 0
+        left_upper_arm_flex_score = 0
+        left_upper_arm_side_score = 0
+        left_upper_arm_shoulder_rise = 0
+
+        right_flexion = self.UA_degrees[0]
+        left_flexion = self.UA_degrees[1]
+
+        right_side = self.UA_degrees[2]
+        left_side = self.UA_degrees[3]
+
+        right_shoulder_rise = self.UA_degrees[4]
+        left_shoulder_rise = self.UA_degrees[5]
+
+        if -20 <= right_flexion < 20:
+            right_upper_arm_reba_score += 1
+            right_upper_arm_flex_score += 1
+        if 20 <= right_flexion < 45:
+            right_upper_arm_reba_score += 2
+            right_upper_arm_flex_score += 2
+        if right_flexion < -20:
+            right_upper_arm_reba_score += 2
+            right_upper_arm_flex_score += 2
+        if 45 <= right_flexion < 90:
+            right_upper_arm_reba_score += 3
+            right_upper_arm_flex_score += 3
+        if 90 <= right_flexion:
+            right_upper_arm_reba_score += 4
+            right_upper_arm_flex_score += 4
+
+        if -20 <= left_flexion < 20:
+            left_upper_arm_reba_score += 1
+            left_upper_arm_flex_score += 1
+        if 20 <= left_flexion < 45:
+            left_upper_arm_reba_score += 2
+            left_upper_arm_flex_score += 2
+        if left_flexion < -20:
+            left_upper_arm_reba_score += 2
+            left_upper_arm_flex_score += 2
+        if 45 <= left_flexion < 90:
+            left_upper_arm_reba_score += 3
+            left_upper_arm_flex_score += 3
+        if 90 <= left_flexion:
+            left_upper_arm_reba_score += 4
+            left_upper_arm_flex_score += 4
+
+            # for side bending
+        if abs(right_side) > 2:
+            right_upper_arm_reba_score += 1
+            right_upper_arm_side_score += 1
+
+        if abs(left_side) > 2:
+            left_upper_arm_reba_score += 1
+            left_upper_arm_side_score += 1
+
+            # for shoulder rise
+
+        if right_shoulder_rise > 90:
+            right_upper_arm_reba_score += 1
+            right_upper_arm_shoulder_rise += 1
+
+        if left_shoulder_rise > 90:
+            left_upper_arm_reba_score += 1
+            left_upper_arm_shoulder_rise += 1
+
+        return [right_upper_arm_reba_score, left_upper_arm_reba_score]

rofunc/utils/ergolab/reba_calculate/pose2degree/Util.py

@@ -0,0 +1,53 @@
+import math
+import numpy as np
+
+
+ # multiplication of two quaternion(representing the orientation of joints)
+def multiply_two_quaternion(q1, q2):
+        a = q1[0]
+        b = q1[1]
+        c = q1[2]
+        d = q1[3]
+        e = q2[0]
+        f = q2[1]
+        g = q2[2]
+        h = q2[3]
+
+        m0 = round(a * e - b * f - c * g - d * h, 1)
+        m1 = round(b * e + a * f + c * h - d * g, 1)
+        m2 = round(a * g - b * h + c * e + d * f, 1)
+        m3 = round(a * h + b * g - c * f + d * e, 1)
+        return [m0, m1, m2, m3]
+
+
+def get_angle_between_degs(v1, v2):
+    len_v1 = math.sqrt(v1[0] ** 2 + v1[1] ** 2 + v1[2] ** 2)
+    len_v2 = math.sqrt(v2[0] ** 2 + v2[1] ** 2 + v2[2] ** 2)
+
+    result = math.acos(round(np.dot(v1, v2) / (len_v1 * len_v2), 3)) * 180 / math.pi
+    return result
+
+
+def get_distance_between(p1, p2):
+    result = [x + y for x, y in zip(p2, np.dot(p1, -1))]
+    return math.sqrt(result[0] ** 2 + result[1] ** 2 + result[2] ** 2)
+
+
+def normalization(vector):
+    l = math.sqrt(vector[0] ** 2 + vector[1] ** 2 + vector[2] ** 2)
+    if l == 0:
+        l += 0.01
+    normal_vector = [vector[0] / l, vector[1] / l, vector[2] / l]
+    return normal_vector
+
+
+def find_rotation_quaternion(outer_quaternion, inner_quaternion):
+    conjucate = [outer_quaternion[0], -outer_quaternion[1], -outer_quaternion[2], -outer_quaternion[3]]
+    length = math.sqrt(outer_quaternion[0] ** 2 + outer_quaternion[1] ** 2 +
+                       outer_quaternion[2] ** 2 + outer_quaternion[3] ** 2)
+    if length ==0:
+        inverse =[0,0,0,-1]
+    else:
+        inverse = np.dot(conjucate, (1 / length))
+    rotation = multiply_two_quaternion(inner_quaternion, inverse)
+    return rotation