-
Notifications
You must be signed in to change notification settings - Fork 0
/
keyword_extraction.py
executable file
·106 lines (94 loc) · 4.83 KB
/
keyword_extraction.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
import numpy as np
import os
import csv
import pickle
import shutil
train_path = "/cvhci/data/activity/NTU_RGBD/train_set.csv"
test_path = "/cvhci/data/activity/NTU_RGBD/test_set.csv"
with open(train_path, 'r') as f:
data_train = list(csv.reader(f, delimiter=";"))
with open(test_path, 'r') as f:
data_test = list(csv.reader(f, delimiter=";"))
# print(data_train[0])
train_key_list = []
test_key_list = []
for index in data_train:
train_key_list.append(index[0].split('/')[-1].split('_rgb')[0])
for index in data_test:
test_key_list.append(index[0].split('/')[-1].split('_rgb')[0])
def find_files(directory, suffix='.png'):
if not os.path.exists(directory):
raise ValueError("Directory not found {}".format(directory))
matches = []
for root, dirnames, filenames in os.walk(directory):
# print(filenames)
for dirname in dirnames:
for root, dirnames, filenames in os.walk(directory + dirname):
# print(dirnames)
for filename in filenames:
full_path = os.path.join(directory + dirname, filename)
# if filename.endswith(suffix):
matches.append(full_path)
return matches
all_list = []
one_shot_train = find_files("/cvhci/data/activity/kpeng/ntu_oneshot/one_shot/train/", '.skeleton.png')
#print(one_shot_train)
one_shot_sample = find_files("/cvhci/data/activity/kpeng/ntu_oneshot/one_shot/samples/", '.skeleton.png')
one_shot_test = find_files("/cvhci/data/activity/kpeng/ntu_oneshot/one_shot/test/", '.skeleton.png')
#all_list = one_shot_train + one_shot_sample + one_shot_test
key_list_train = [index.split('.skeleton')[0].split('/')[-1] for index in one_shot_train]
label_list_train = [index.split('/')[-2] for index in one_shot_train]
key_list_test = [index.split('.skeleton')[0].split('/')[-1] for index in one_shot_test]
label_list_test = [index.split('/')[-2] for index in one_shot_test]
key_list_sample = [index.split('.skeleton')[0].split('/')[-1] for index in one_shot_sample]
label_list_sample = [index.split('/')[-2] for index in one_shot_sample]
print(key_list_train)
print(label_list_train)
for sample in train_key_list:
if sample in key_list_train:
save_path = "/cvhci/data/activity/kpeng/ntu_train_test/train/"
position = key_list_train.index(sample)
os.makedirs(save_path + label_list_train[position], exist_ok=True)
save_path = save_path + label_list_train[position] + '/' + sample + '.skeleton.png'
data_path = one_shot_train[position]
shutil.copyfile(data_path, save_path)
if sample in key_list_test:
save_path = "/cvhci/data/activity/kpeng/ntu_train_test/train/"
position = key_list_test.index(sample)
os.makedirs(save_path + label_list_test[position], exist_ok=True)
save_path = save_path + label_list_test[position] + '/' + sample + '.skeleton.png'
data_path = one_shot_test[position]
shutil.copyfile(data_path, save_path)
if sample in key_list_sample:
save_path = "/cvhci/data/activity/kpeng/ntu_train_test/train/"
position = key_list_sample.index(sample)
os.makedirs(save_path + label_list_sample[position], exist_ok=True)
save_path = save_path + label_list_sample[position] + '/' + sample + '.skeleton.png'
data_path = one_shot_sample[position]
shutil.copyfile(data_path, save_path)
for sample in test_key_list:
if sample in key_list_train:
save_path = "/cvhci/data/activity/kpeng/ntu_train_test/test/"
position = key_list_train.index(sample)
os.makedirs(save_path + label_list_train[position], exist_ok=True)
save_path = save_path + label_list_train[position] + '/' + sample + '.skeleton.png'
data_path = one_shot_train[position]
shutil.copyfile(data_path, save_path)
if sample in key_list_test:
save_path = "/cvhci/data/activity/kpeng/ntu_train_test/test/"
position = key_list_test.index(sample)
os.makedirs(save_path + label_list_test[position], exist_ok=True)
save_path = save_path + label_list_test[position] + '/' + sample + '.skeleton.png'
data_path = one_shot_test[position]
shutil.copyfile(data_path, save_path)
if sample in key_list_sample:
save_path = "/cvhci/data/activity/kpeng/ntu_train_test/test/"
position = key_list_sample.index(sample)
os.makedirs(save_path + label_list_sample[position], exist_ok=True)
save_path = save_path + label_list_sample[position] + '/' + sample + '.skeleton.png'
data_path = one_shot_sample[position]
shutil.copyfile(data_path, save_path)
# with open('ntu_train.pkl', 'wb') as f:
# pickle.dump(train_key_list, f)
# with open('ntu_test.pkl', 'wb') as f:
# pickle.dump(test_key_list, f)