Skip to content
Snippets Groups Projects
Select Git revision
  • dev default
  • dataset
  • jean-zay
  • main protected
4 results
Blame History Permalink
pose.py 2.25 KiB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67 






import math
import numpy as np
import json
from scipy.spatial.transform import Rotation


def convert2(xyz):
    (R, P, Y) = (xyz[0], xyz[1], xyz[2])
    Q = Rotation.from_euler(seq='xyz', angles=[R, P, Y], degrees=False).as_quat()
    r = Rotation.from_quat(Q)
    rotation = r.as_matrix()

    return rotation


def transform_pose(data_name):
    transformation = np.matrix([[0.0000000, -1.0000000, 0.0000000],
                                [0.0000000, 0.0000000, -1.0000000],
                                [1.0000000, 0.0000000, 0.0000000]])
    cont1, cont2, cont3 = 0, 0, 0
    for p in range(4995):
        with open(f'{data_name}/Pose/{p}.json', 'r') as f:
          data = json.load(f)

        for i in range(len(data)):
            if data[i]['id'] == 4:
                cont1 += 1
                rpy = data[i]['pose']['rpy']
                rot = convert2(rpy)
                R_exp = transformation @ rot
                R_exp = np.array(R_exp)

                xyz = data[i]['pose']['xyz']
                T_exp = transformation @ xyz
                T_exp = np.array(T_exp)
                num_arr = np.c_[R_exp, T_exp[0]]
                np.save(f'{data_name}/Pose_transformed/Banana/{p}.npy', num_arr)  # save

            elif data[i]['id'] == 5:
                cont2 += 1
                rpy = data[i]['pose']['rpy']
                rot = convert2(rpy)
                R_exp = transformation @ rot
                R_exp = np.array(R_exp)

                xyz = data[i]['pose']['xyz']
                T_exp = transformation @ xyz
                T_exp = np.array(T_exp)
                num_arr = np.c_[R_exp, T_exp[0]]
                np.save(f'{data_name}/Pose_transformed/Orange/{p}.npy', num_arr)  # save
            elif data[i]['id'] == 6:
                cont3 += 1
                rpy = data[i]['pose']['rpy']
                rot = convert2(rpy)
                R_exp = transformation @ rot
                R_exp = np.array(R_exp)

                xyz = data[i]['pose']['xyz']
                T_exp = transformation @ xyz
                T_exp = np.array(T_exp)
                num_arr = np.c_[R_exp, T_exp[0]]
                np.save(f'{data_name}/Pose_transformed/Pear/{p}.npy', num_arr)  # save
            else:
                continue
    print(cont1, cont2, cont3)