-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathconvert_from_imagenet_to_coco.py
67 lines (56 loc) · 2.02 KB
/
convert_from_imagenet_to_coco.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
import json
import os
import sys
from six.moves import urllib
#get synset_to_human dict
print('Loading synset id to human readable class name dictionary...')
base_url = 'https://raw.githubusercontent.com/tensorflow/models/master/research/inception/inception/data/'
synset_to_human_url = '{}/imagenet_metadata.txt'.format(base_url)
filename, _ = urllib.request.urlretrieve(synset_to_human_url)
synset_to_human_list = open(filename).readlines()
synset_to_human = {}
for s in synset_to_human_list:
parts = s.strip().split('\t')
assert len(parts) == 2
synset = parts[0]
human = parts[1]
synset_to_human[synset] = human
#get images and classes (training only)
print('Reading ground truth from file structure and creating COCO database...')
base_folder = '/home/ubuntu/imagenet_data/'
split = 'train'
folder = base_folder + split + '/'
files = os.listdir(folder)
images = []
annotations = []
cats = []
ann_id_count = 0
for directory in files:
if os.path.isdir(folder+directory):
cat_id = directory
cats.append(cat_id)
path = os.listdir(folder+directory+'/')
for filename in path:
im = {}
im['file_name'] = folder+directory+'/'+filename
im['id'] = filename.split('.')[0]
images.append(im)
ann = {}
ann['image_id'] = im['id']
ann['id'] = str(ann_id_count)
ann_id_count += 1
ann['category_id'] = cat_id
annotations.append(ann)
print(len(images),len(annotations),len(cats))
categories = []
for cat_id in cats:
cat = {}
cat['id'] = cat_id
cat['name'] = synset_to_human[cat_id]
categories.append(cat)
data = {}
data['images'] = images
data['categories'] = categories
data['annotations'] = annotations
print('Saving COCO training database...')
json.dump(data,open('ILSVRC2014/ILSVRC2014_'+split+'_classification.json','w'))