Pytorch-template/generate_txt.py at main · Zehui-Lin/Pytorch-template · 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
import os
from utils import check_dir

data_path = './'
txtdir_path = './data/txt'
class1_path = data_path + ''
class2_path = data_path + ''

check_dir(txtdir_path)

class1_list = os.listdir(class1_path)
class2_list = os.listdir(class2_path)

class1_num = len(class1_list)
class2_num = len(class2_list)


train_file = open(os.path.join(txtdir_path, "train.txt"), "w")
val_file = open(os.path.join(txtdir_path, "val.txt"), "w")
test_file = open(os.path.join(txtdir_path, "test.txt"), "w")


# train
for data in class1_list[0:int(0.6*class1_num)]:
    data_name = os.path.join(class1_path, data)
    train_file.write("{},{}\n".format(data_name, 1))

for data in class2_list[0:int(0.6*class2_num)]:
    data_name = os.path.join(class2_path, data)
    train_file.write("{},{}\n".format(data_name, 0))

# val
for data in class1_list[int(0.6*class1_num):int(0.8*class1_num)]:
    data_name = os.path.join(class1_path, data)
    val_file.write("{},{}\n".format(data_name, 1))

for data in class2_list[int(0.6*class2_num):int(0.8*class2_num)]:
    data_name = os.path.join(class2_path, data)
    val_file.write("{},{}\n".format(data_name, 0))

# test
for data in class1_list[int(0.8*class1_num):]:
    data_name = os.path.join(class1_path, data)
    test_file.write("{},{}\n".format(data_name, 1))

for data in class2_list[int(0.8*class2_num):]:
    data_name = os.path.join(class2_path, data)
    test_file.write("{},{}\n".format(data_name, 0))


train_file.close()
val_file.close()
test_file.close()

print("generate  finished!")


train_file = open(os.path.join(txtdir_path, "train.txt"), "r")
val_file = open(os.path.join(txtdir_path, "val.txt"), "r")
test_file = open(os.path.join(txtdir_path, "test.txt"), "r")
all_file = open(os.path.join(txtdir_path, "all.txt"), "w")
all_file.writelines([*train_file.readlines(), *val_file.readlines(), *test_file.readlines()])
all_file.close()