-
Notifications
You must be signed in to change notification settings - Fork 0
/
test_split.py
50 lines (43 loc) · 1.47 KB
/
test_split.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
"""
@authors: Zhongjie Ye, Dongchao Yang, Helin Wang
@Introduction: Randomly Select Test data from Training data.
"""
import os
import random
import shutil
import argparse
def main(args):
data_path = args.data_path
ls = os.listdir(os.path.join(data_path, 'train'))
if not os.path.isdir(os.path.join(data_path, 'test')):
os.mkdir(os.path.join(data_path, 'test'))
one_num = args.test_sample_num // 2
zero_num = args.test_sample_num - one_num
one_ls = []
zero_ls = []
for a in ls:
label = int(a[-5])
if label == 0:
zero_num += 1
zero_ls.append(a)
else:
one_num += 1
one_ls.append(a)
print(len(zero_ls))
print(len(one_ls))
tmp_one = random.sample(one_ls, one_num)
tmp_zero = random.sample(zero_ls, zero_num)
for a in tmp_one:
full_path = os.path.join(data_path, 'train', a)
despath = os.path.join(data_path, 'test', a)
shutil.move(full_path, despath)
for a in tmp_zero:
full_path = os.path.join(data_path, 'train', a)
despath = os.path.join(data_path, 'test', a)
shutil.move(full_path, despath)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description='Randomly Select Test data from Training data.')
parser.add_argument('--data_path', type=str) # /home/pkusz/home/PKU_team/new_data/
parser.add_argument('--test_sample_num', type=int, default=800)
args = parser.parse_args()
main(args)