taskswithcode/DeticChatGPT

1# Copyright (c) Facebook, Inc. and its affiliates.2import argparse3import json4import os5import cv26from nltk.corpus import wordnet7 8if __name__ == '__main__':9    parser = argparse.ArgumentParser()10    parser.add_argument('--imagenet_path', default='datasets/imagenet/ImageNet-LVIS')11    parser.add_argument('--lvis_meta_path', default='datasets/lvis/lvis_v1_val.json')12    parser.add_argument('--out_path', default='datasets/imagenet/annotations/imagenet_lvis_image_info.json')13    args = parser.parse_args()14 15    print('Loading LVIS meta')16    data = json.load(open(args.lvis_meta_path, 'r'))17    print('Done')18    synset2cat = {x['synset']: x for x in data['categories']}19    count = 020    images = []21    image_counts = {}22    folders = sorted(os.listdir(args.imagenet_path))23    for i, folder in enumerate(folders):24        class_path = args.imagenet_path + folder25        files = sorted(os.listdir(class_path))26        synset = wordnet.synset_from_pos_and_offset('n', int(folder[1:])).name()27        cat = synset2cat[synset]28        cat_id = cat['id']29        cat_name = cat['name']30        cat_images = []31        for file in files:32            count = count + 133            file_name = '{}/{}'.format(folder, file)34            img = cv2.imread('{}/{}'.format(args.imagenet_path, file_name))35            h, w = img.shape[:2]36            image = {37                'id': count,38                'file_name': file_name,39                'pos_category_ids': [cat_id],40                'width': w,41                'height': h42            }43            cat_images.append(image)44        images.extend(cat_images)45        image_counts[cat_id] = len(cat_images)46        print(i, cat_name, len(cat_images))47    print('# Images', len(images))48    for x in data['categories']:49        x['image_count'] = image_counts[x['id']] if x['id'] in image_counts else 050    out = {'categories': data['categories'], 'images': images, 'annotations': []}51    print('Writing to', args.out_path)52    json.dump(out, open(args.out_path, 'w'))53