-
Notifications
You must be signed in to change notification settings - Fork 235
/
img_generator.py
51 lines (38 loc) · 1.29 KB
/
img_generator.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
#!/usr/bin/env python
# -- coding: utf-8 --
"""
Copyright (c) 2019. All rights reserved.
Created by C. L. Wang on 2019/1/3
"""
from project_utils import *
from root_dir import ROOT_DIR
from yolo3_predict import YOLO
def generate_data(file_name):
print('执行开始')
out_file = os.path.join(ROOT_DIR, 'dataset', file_name)
create_file(out_file)
up_folder = os.path.abspath(os.path.join(ROOT_DIR, '..'))
data_folder = os.path.join(up_folder, 'data_set', 'XX-Images-57w-416')
path_list, name_list = traverse_dir_files(data_folder)
yolo = YOLO()
for count, (path, name) in enumerate(zip(path_list, name_list)):
# print(path, name)
try:
objects_line = yolo.detect_objects_of_image(path)
except Exception as e:
print("错误: {}".format(name))
continue
if not objects_line:
continue
# print(objects_line)
print("已处理: {}".format(name))
data_line = '{}---{}'.format(name, objects_line)
write_line_utf8(out_file, data_line)
if count % 1000 == 0:
print('count: {}'.format(count))
# break
yolo.close_session()
print('执行结束')
if __name__ == '__main__':
file_name = 'dataset_57w_objects.txt'
generate_data(file_name)