forked from yysijie/st-gcn
-
Notifications
You must be signed in to change notification settings - Fork 0
/
ntu_gendata.py
122 lines (103 loc) · 3.61 KB
/
ntu_gendata.py
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
import os
import sys
import pickle
import argparse
import numpy as np
from numpy.lib.format import open_memmap
from utils.ntu_read_skeleton import read_xyz
training_subjects = [
1, 2, 4, 5, 8, 9, 13, 14, 15, 16, 17, 18, 19, 25, 27, 28, 31, 34, 35, 38
]
training_cameras = [2, 3]
max_body = 2
num_joint = 25
max_frame = 300
toolbar_width = 30
def print_toolbar(rate, annotation=''):
# setup toolbar
sys.stdout.write("{}[".format(annotation))
for i in range(toolbar_width):
if i * 1.0 / toolbar_width > rate:
sys.stdout.write(' ')
else:
sys.stdout.write('-')
sys.stdout.flush()
sys.stdout.write(']\r')
def end_toolbar():
sys.stdout.write("\n")
def gendata(data_path,
out_path,
ignored_sample_path=None,
benchmark='xview',
part='eval'):
if ignored_sample_path != None:
with open(ignored_sample_path, 'r') as f:
ignored_samples = [
line.strip() + '.skeleton' for line in f.readlines()
]
else:
ignored_samples = []
sample_name = []
sample_label = []
for filename in os.listdir(data_path):
if filename in ignored_samples:
continue
action_class = int(
filename[filename.find('A') + 1:filename.find('A') + 4])
subject_id = int(
filename[filename.find('P') + 1:filename.find('P') + 4])
camera_id = int(
filename[filename.find('C') + 1:filename.find('C') + 4])
if benchmark == 'xview':
istraining = (camera_id in training_cameras)
elif benchmark == 'xsub':
istraining = (subject_id in training_subjects)
else:
raise ValueError()
if part == 'train':
issample = istraining
elif part == 'val':
issample = not (istraining)
else:
raise ValueError()
if issample:
sample_name.append(filename)
sample_label.append(action_class - 1)
with open('{}/{}_label.pkl'.format(out_path, part), 'wb') as f:
pickle.dump((sample_name, list(sample_label)), f)
# np.save('{}/{}_label.npy'.format(out_path, part), sample_label)
fp = open_memmap(
'{}/{}_data.npy'.format(out_path, part),
dtype='float32',
mode='w+',
shape=(len(sample_label), 3, max_frame, num_joint, max_body))
for i, s in enumerate(sample_name):
print_toolbar(i * 1.0 / len(sample_label),
'({:>5}/{:<5}) Processing {:>5}-{:<5} data: '.format(
i + 1, len(sample_name), benchmark, part))
data = read_xyz(
os.path.join(data_path, s), max_body=max_body, num_joint=num_joint)
fp[i, :, 0:data.shape[1], :, :] = data
end_toolbar()
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='NTU-RGB-D Data Converter.')
parser.add_argument(
'--data_path', default='data/NTU-RGB-D/nturgb+d_skeletons')
parser.add_argument(
'--ignored_sample_path',
default='resource/NTU-RGB-D/samples_with_missing_skeletons.txt')
parser.add_argument('--out_folder', default='data/NTU-RGB-D')
benchmark = ['xsub', 'xview']
part = ['train', 'val']
arg = parser.parse_args()
for b in benchmark:
for p in part:
out_path = os.path.join(arg.out_folder, b)
if not os.path.exists(out_path):
os.makedirs(out_path)
gendata(
arg.data_path,
out_path,
arg.ignored_sample_path,
benchmark=b,
part=p)