diff --git a/examples/learning_ml/example_felt.py b/examples/learning_ml/example_felt.py index 23b488f56..94f5535f0 100644 --- a/examples/learning_ml/example_felt.py +++ b/examples/learning_ml/example_felt.py @@ -8,14 +8,31 @@ import numpy as np import os import rofunc as rf +import pandas as pd # --- Data processing --- def data_process(data_dir): all_files = rf.file.list_absl_path(data_dir, recursive=False, prefix='trial') for file in all_files: - hand_rigid = np.load(os.path.join(file, 'mocap_hand_rigid.npy')) - object_rigid = np.load(os.path.join(file, 'mocap_object_rigid.npy')) + hand_rigid = pd.read_csv(os.path.join(file, 'mocap_hand_rigid.csv')) + object_rigid = pd.read_csv(os.path.join(file, 'mocap_object_rigid.csv')) + hand_marker_positions = pd.read_csv(os.path.join(file, 'mocap_hand.csv')) + object_marker_positions = pd.read_csv(os.path.join(file, 'mocap_object.csv')) + + def get_center_position(df): + data = df.to_numpy().reshape((len(df.to_numpy()), -1, 3)) + return np.mean(data, axis=1) + + def get_orientation(df): + data = df.to_numpy().reshape((len(df.to_numpy()), 3, 3)) + data = np.array([rf.robolab.quaternion_from_matrix(rf.robolab.homo_matrix_from_rot_matrix(i)) for i in data]) + return data + + hand_position = get_center_position(hand_marker_positions) + object_position = get_center_position(object_marker_positions) + hand_ori = get_orientation(hand_rigid) + object_ori = get_orientation(object_rigid) return demos_x diff --git a/rofunc/utils/robolab/coord/transform.py b/rofunc/utils/robolab/coord/transform.py index a5e614ab6..998ece1d2 100644 --- a/rofunc/utils/robolab/coord/transform.py +++ b/rofunc/utils/robolab/coord/transform.py @@ -948,6 +948,23 @@ def rot_matrix_from_euler(ai, aj, ak, axes='sxyz'): return rot_matrix +def homo_matrix_from_rot_matrix(rot_matrix, translation=None): + """Construct homogeneous matrix from rotation matrix + >>> R = homo_matrix_from_rot_matrix(np.array([[-0.64287284, 0.05872664, -0.76371834], \ + [-0.68124272, 0.41198225, 0.60512725], \ + [ 0.35017549, 0.90929743, -0.2248451 ]])) + + :param rot_matrix: [3, 3] array + :param translation: quaternion + """ + check_rot_matrix(rot_matrix) + homo_matrix = np.identity(4) + homo_matrix[:3, :3] = rot_matrix + if translation is not None: + homo_matrix[:, 3] = translation + return homo_matrix + + def euler_from_matrix(matrix, axes='sxyz'): """Return Euler angles from rotation matrix for specified axis sequence. axes : One of 24 axis sequences as string or encoded tuple @@ -1080,35 +1097,6 @@ def quaternion_about_axis(angle, axis): return quaternion -def homo_matrix_from_quaternion(quaternion): - """Return homogeneous rotation matrix from quaternion. - >>> R = homo_matrix_from_quaternion([0.06146124, 0, 0, 0.99810947]) - >>> np.allclose(R, rotation_matrix(0.123, (1, 0, 0))) - True - """ - q = np.array(quaternion[:4], dtype=np.float64, copy=True) - nq = np.dot(q, q) - if nq < _EPS: - return np.identity(4) - q *= math.sqrt(2.0 / nq) - q = np.outer(q, q) - return np.array(( - (1.0 - q[1, 1] - q[2, 2], q[0, 1] - q[2, 3], q[0, 2] + q[1, 3], 0.0), - (q[0, 1] + q[2, 3], 1.0 - q[0, 0] - q[2, 2], q[1, 2] - q[0, 3], 0.0), - (q[0, 2] - q[1, 3], q[1, 2] + q[0, 3], 1.0 - q[0, 0] - q[1, 1], 0.0), - (0.0, 0.0, 0.0, 1.0) - ), dtype=np.float64) - - -def rot_matrix_from_quaternion(quaternion): - """Return rotation matrix from quaternion. - >>> R = rot_matrix_from_quaternion([0.06146124, 0, 0, 0.99810947]) - """ - homo_matrix = homo_matrix_from_quaternion(quaternion) - rot_matrix = homo_matrix[:3, :3] - return rot_matrix - - def quaternion_from_matrix(matrix): """Return quaternion from rotation matrix. >>> R = rotation_matrix(0.123, (1, 2, 3)) @@ -1241,6 +1229,35 @@ def random_quaternion(rand=None): np.cos(t2) * r2), dtype=np.float64) +def homo_matrix_from_quaternion(quaternion): + """Return homogeneous rotation matrix from quaternion. + >>> R = homo_matrix_from_quaternion([0.06146124, 0, 0, 0.99810947]) + >>> np.allclose(R, rotation_matrix(0.123, (1, 0, 0))) + True + """ + q = np.array(quaternion[:4], dtype=np.float64, copy=True) + nq = np.dot(q, q) + if nq < _EPS: + return np.identity(4) + q *= math.sqrt(2.0 / nq) + q = np.outer(q, q) + return np.array(( + (1.0 - q[1, 1] - q[2, 2], q[0, 1] - q[2, 3], q[0, 2] + q[1, 3], 0.0), + (q[0, 1] + q[2, 3], 1.0 - q[0, 0] - q[2, 2], q[1, 2] - q[0, 3], 0.0), + (q[0, 2] - q[1, 3], q[1, 2] + q[0, 3], 1.0 - q[0, 0] - q[1, 1], 0.0), + (0.0, 0.0, 0.0, 1.0) + ), dtype=np.float64) + + +def rot_matrix_from_quaternion(quaternion): + """Return rotation matrix from quaternion. + >>> R = rot_matrix_from_quaternion([0.06146124, 0, 0, 0.99810947]) + """ + homo_matrix = homo_matrix_from_quaternion(quaternion) + rot_matrix = homo_matrix[:3, :3] + return rot_matrix + + def random_rotation_matrix(rand=None): """Return uniform random rotation matrix. rnd: array like