-
Notifications
You must be signed in to change notification settings - Fork 13
/
Copy pathprepare_coco_data.py
71 lines (48 loc) · 1.55 KB
/
prepare_coco_data.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
import numpy as np
import os
from lib.dataset.coco_data import BoxInfo
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--mscocodir', type=str,default='../pubdata/mscoco', help='detect with coco or face',required=False)
args = parser.parse_args()
coco_dir=args.mscocodir
train_im_path = os.path.join(coco_dir,'train2017')
train_ann_path = os.path.join(coco_dir,'annotations/instances_train2017.json')
val_im_path = os.path.join(coco_dir,'val2017')
val_ann_path = os.path.join(coco_dir,'annotations/instances_val2017.json')
train_data=BoxInfo(train_im_path,train_ann_path)
fw = open('train.txt', 'w')
for meta in train_data.metas:
fname, boxes = meta.img_url, meta.bbox
tmp_str = ''
tmp_str =tmp_str+ fname+'|'
for box in boxes:
data = ' %d,%d,%d,%d,%d'%(box[0], box[1], box[2], box[3],box[4])
tmp_str=tmp_str+data
if len(boxes) == 0:
print(tmp_str)
continue
####err box?
if box[2] <= 0 or box[3] <= 0:
pass
else:
fw.write(tmp_str + '\n')
fw.close()
val_data=BoxInfo(val_im_path,val_ann_path)
fw = open('val.txt', 'w')
for meta in val_data.metas:
fname, boxes = meta.img_url, meta.bbox
tmp_str = ''
tmp_str = tmp_str + fname + '|'
for box in boxes:
data = ' %d,%d,%d,%d,%d' % (box[0], box[1], box[2], box[3], box[4])
tmp_str = tmp_str + data
if len(boxes) == 0:
print(tmp_str)
continue
####err box?
if box[2] <= 0 or box[3] <= 0:
pass
else:
fw.write(tmp_str + '\n')
fw.close()