Spaces:
Paused
Paused
| # Copyright (c) Facebook, Inc. and its affiliates. | |
| import argparse | |
| import copy | |
| import json | |
| if __name__ == '__main__': | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument('--imagenet_path', default='datasets/imagenet/annotations/imagenet-21k_image_info.json') | |
| parser.add_argument('--lvis_path', default='datasets/lvis/lvis_v1_train.json') | |
| parser.add_argument('--save_categories', default='') | |
| parser.add_argument('--not_save_imagenet', action='store_true') | |
| parser.add_argument('--not_save_lvis', action='store_true') | |
| parser.add_argument('--mark', default='lvis-21k') | |
| args = parser.parse_args() | |
| print('Loading', args.imagenet_path) | |
| in_data = json.load(open(args.imagenet_path, 'r')) | |
| print('Loading', args.lvis_path) | |
| lvis_data = json.load(open(args.lvis_path, 'r')) | |
| categories = copy.deepcopy(lvis_data['categories']) | |
| cat_count = max(x['id'] for x in categories) | |
| synset2id = {x['synset']: x['id'] for x in categories} | |
| name2id = {x['name']: x['id'] for x in categories} | |
| in_id_map = {} | |
| for x in in_data['categories']: | |
| if x['synset'] in synset2id: | |
| in_id_map[x['id']] = synset2id[x['synset']] | |
| elif x['name'] in name2id: | |
| in_id_map[x['id']] = name2id[x['name']] | |
| x['id'] = name2id[x['name']] | |
| else: | |
| cat_count = cat_count + 1 | |
| name2id[x['name']] = cat_count | |
| in_id_map[x['id']] = cat_count | |
| x['id'] = cat_count | |
| categories.append(x) | |
| print('lvis cats', len(lvis_data['categories'])) | |
| print('imagenet cats', len(in_data['categories'])) | |
| print('merge cats', len(categories)) | |
| filtered_images = [] | |
| for x in in_data['images']: | |
| x['pos_category_ids'] = [in_id_map[xx] for xx in x['pos_category_ids']] | |
| x['pos_category_ids'] = [xx for xx in \ | |
| sorted(set(x['pos_category_ids'])) if xx >= 0] | |
| if len(x['pos_category_ids']) > 0: | |
| filtered_images.append(x) | |
| in_data['categories'] = categories | |
| lvis_data['categories'] = categories | |
| if not args.not_save_imagenet: | |
| in_out_path = args.imagenet_path[:-5] + '_{}.json'.format(args.mark) | |
| for k, v in in_data.items(): | |
| print('imagenet', k, len(v)) | |
| print('Saving Imagenet to', in_out_path) | |
| json.dump(in_data, open(in_out_path, 'w')) | |
| if not args.not_save_lvis: | |
| lvis_out_path = args.lvis_path[:-5] + '_{}.json'.format(args.mark) | |
| for k, v in lvis_data.items(): | |
| print('lvis', k, len(v)) | |
| print('Saving LVIS to', lvis_out_path) | |
| json.dump(lvis_data, open(lvis_out_path, 'w')) | |
| if args.save_categories != '': | |
| for x in categories: | |
| for k in ['image_count', 'instance_count', 'synonyms', 'def']: | |
| if k in x: | |
| del x[k] | |
| CATEGORIES = repr(categories) + " # noqa" | |
| open(args.save_categories, 'wt').write(f"CATEGORIES = {CATEGORIES}") | |