MetaGenDataUtils/process_neos_data_sandbox.py at main · MetaGenAI/MetaGenDataUtils · GitHub

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
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138


fileName = "data/example/1/ID2B00_streams.dat"
file = open(fileName, mode='rb')
fileContent = file.read()

oldToNewNodeInts = {
1:1,
9:10,
16:17,
45:46,
}

#%%
import struct
import numpy as np

def read_float(data_iter):
    bs = [next(data_iter,None) for _ in range(4)]
    if np.any(np.array(bs) == None):
        return None
    bs = bytes(bs)
    # bytes = b''.join(bytes)
    x, = struct.unpack('f',bs)
    return x

def read_int(data_iter):
    bs = [next(data_iter,None) for _ in range(4)]
    if np.any(np.array(bs) == None):
        return None
    bs = bytes(bs)
    # bytes = b''.join(bytes)
    x, = struct.unpack('i',bs)
    return x

def read_bool(data_iter):
    bs = [next(data_iter,None) for _ in range(1)]
    if np.any(np.array(bs) == None):
        return None
    bs = bytes(bs)
    # bytes = b''.join(bytes)
    x, = struct.unpack('?',bs)
    return x

#%%
data_iter = iter(fileContent)

######HEADING
# READ absolute time
absolute_time = read_float(data_iter)

# READ version number
version_number = read_int(data_iter)
print("VERSION NUMBER", version_number)

num_body_nodes = version_number
if version_number >= 1000:
    #READ relative avatar scale
    rel_avatar_scale_x = read_float(data_iter)
    rel_avatar_scale_y = read_float(data_iter)
    rel_avatar_scale_z = read_float(data_iter)
    #READ number of body nodes
    num_body_nodes = read_int(data_iter)


body_node_data = {}
for i in range(num_body_nodes):
    #READ body node type
    nodeInt = read_int(data_iter)
    if version_number == 1000:
        nodeInt = oldToNewNodeInts[nodeInt]
    body_node_data[nodeInt] = {}
    print(nodeInt)
    #READ if scale stream exists
    scale_exists = read_bool(data_iter)
    body_node_data[nodeInt]["scale_exists"] = scale_exists
    if scale_exists:
        body_node_data[nodeInt]["scale_stream"] = []
    #READ if position stream exists
    pos_exists = read_bool(data_iter)
    body_node_data[nodeInt]["pos_exists"] = pos_exists
    if pos_exists:
        body_node_data[nodeInt]["pos_stream"] = []
    #READ if rotation stream exists
    rot_exists = read_bool(data_iter)
    body_node_data[nodeInt]["rot_exists"] = rot_exists
    if rot_exists:
        body_node_data[nodeInt]["rot_stream"] = []

#READ whether hands are tracked
hands_are_tracked = read_bool(data_iter)
#READ whether metacarpals are tracked
metacarpals_are_tracked = read_bool(data_iter)

#######STREAM

while True:
    #READ deltaT
    deltaT = read_float(data_iter)
    if deltaT is None:
        break
    #READ bodyNode transforms
    for key, bodyNode in body_node_data.items():
        if bodyNode["scale_exists"]:
            scalex = read_float(data_iter)
            scaley = read_float(data_iter)
            scalez = read_float(data_iter)
            scale = [scalex,scaley,scalez]
            bodyNode["scale_stream"] += [scale]
        if bodyNode["pos_exists"]:
            posx = read_float(data_iter)
            posy = read_float(data_iter)
            posz = read_float(data_iter)
            pos = [posx,posy,posz]
            bodyNode["pos_stream"] += [pos]
        if bodyNode["rot_exists"]:
            rotx = read_float(data_iter)
            roty = read_float(data_iter)
            rotz = read_float(data_iter)
            rotw = read_float(data_iter)
            rot = [rotx,roty,rotz,rotw]
            bodyNode["rot_stream"] += [rot]
    #READ finger poses
    if hands_are_tracked:
        #Left hand
        for i in range(23):
            was_succesful = read_bool(data_iter)
            x = read_float(data_iter)
            y = read_float(data_iter)
            z = read_float(data_iter)
            w = read_float(data_iter)
        #Right hand
        for i in range(23):
            was_succesful = read_bool(data_iter)
            x = read_float(data_iter)
            y = read_float(data_iter)
            z = read_float(data_iter)
            w = read_float(data_iter)