Spaces:
Build error
Build error
File size: 2,019 Bytes
87e5035 |
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 |
# Copyright (c) Facebook, Inc. and its affiliates.
import argparse
import json
import os
import cv2
from nltk.corpus import wordnet
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--imagenet_path', default='datasets/imagenet/ImageNet-LVIS')
parser.add_argument('--lvis_meta_path', default='datasets/lvis/lvis_v1_val.json')
parser.add_argument('--out_path', default='datasets/imagenet/annotations/imagenet_lvis_image_info.json')
args = parser.parse_args()
print('Loading LVIS meta')
data = json.load(open(args.lvis_meta_path, 'r'))
print('Done')
synset2cat = {x['synset']: x for x in data['categories']}
count = 0
images = []
image_counts = {}
folders = sorted(os.listdir(args.imagenet_path))
for i, folder in enumerate(folders):
class_path = args.imagenet_path + folder
files = sorted(os.listdir(class_path))
synset = wordnet.synset_from_pos_and_offset('n', int(folder[1:])).name()
cat = synset2cat[synset]
cat_id = cat['id']
cat_name = cat['name']
cat_images = []
for file in files:
count = count + 1
file_name = '{}/{}'.format(folder, file)
img = cv2.imread('{}/{}'.format(args.imagenet_path, file_name))
h, w = img.shape[:2]
image = {
'id': count,
'file_name': file_name,
'pos_category_ids': [cat_id],
'width': w,
'height': h
}
cat_images.append(image)
images.extend(cat_images)
image_counts[cat_id] = len(cat_images)
print(i, cat_name, len(cat_images))
print('# Images', len(images))
for x in data['categories']:
x['image_count'] = image_counts[x['id']] if x['id'] in image_counts else 0
out = {'categories': data['categories'], 'images': images, 'annotations': []}
print('Writing to', args.out_path)
json.dump(out, open(args.out_path, 'w'))
|