# 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'))