-
Notifications
You must be signed in to change notification settings - Fork 0
/
utils.py
59 lines (48 loc) · 1.97 KB
/
utils.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
import os
import time
import random
import shutil
import argparse
def gen_files(args):
root = args.data_path
files_list = []
for roots, _, files in os.walk(root):
for file in files:
if file.endswith(args.postfix):
files_list.append(os.path.join(roots, file))
len_files = len(files_list)
print("数据总数量: ", len_files)
return files_list
def gen_txt(args):
files = gen_files(args)
random.shuffle(files)
files = [file + "\n" for file in files]
if args.val:
trains = files[:int(len(files) * args.ratio)]
vals = files[int(len(files) * args.ratio):]
# vals = files[int(len(files) * args.ratio):int(len(files) * args.ratio*2)]
# test = files[int(len(files) * args.ratio*2):]
else:
trains = files
vals = files
# test = files
if not os.path.exists(args.save_path):
os.mkdir(args.save_path)
with open(args.save_path + r'/train.txt', 'w') as f:
f.writelines(trains)
with open(args.save_path + r'/val.txt', 'w') as f:
f.writelines(vals)
# with open(args.save_path + r'/test.txt', 'w') as f:
# f.writelines(test)
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='文件处理')
parser.add_argument('--data_path', type=str, help='文件路径', default=r"data/drinks_dataset/") # need to do
parser.add_argument('--save_path', type=str, help='保存的路径', default=r"data/drinks_dataset/") # need to do
parser.add_argument('--postfix', type=str, help='要操作的文件后缀名', default=('.jpg','.jpeg'))
parser.add_argument('--val', type=bool, help='是否要切割数据集', default=True)
parser.add_argument('--ratio', type=str, help='训练集占比', default=0.4)
args = parser.parse_args()
print(type(args))
start = time.time()
gen_txt(args)
print(f"Operation finished in {time.time()-start}")