taskswithcode/DeticChatGPT

1# Copyright (c) Facebook, Inc. and its affiliates.2import argparse3import json4 5if __name__ == '__main__':6    parser = argparse.ArgumentParser()7    parser.add_argument("--ann", default='datasets/lvis/lvis_v1_train.json')8    parser.add_argument("--add_freq", action='store_true')9    parser.add_argument("--r_thresh", type=int, default=10)10    parser.add_argument("--c_thresh", type=int, default=100)11    args = parser.parse_args()12 13    print('Loading', args.ann)14    data = json.load(open(args.ann, 'r'))15    cats = data['categories']16    image_count = {x['id']: set() for x in cats}17    ann_count = {x['id']: 0 for x in cats}18    for x in data['annotations']:19        image_count[x['category_id']].add(x['image_id'])20        ann_count[x['category_id']] += 121    num_freqs = {x: 0 for x in ['r', 'f', 'c']}22    for x in cats:23        x['image_count'] = len(image_count[x['id']])24        x['instance_count'] = ann_count[x['id']]25        if args.add_freq:26            freq = 'f'27            if x['image_count'] < args.c_thresh:28                freq = 'c'29            if x['image_count'] < args.r_thresh:30                freq = 'r'31            x['frequency'] = freq32            num_freqs[freq] += 133    print(cats)34    image_counts = sorted([x['image_count'] for x in cats])35    # print('image count', image_counts)36    # import pdb; pdb.set_trace()37    if args.add_freq:38        for x in ['r', 'c', 'f']:39            print(x, num_freqs[x])40    out = cats # {'categories': cats}41    out_path = args.ann[:-5] + '_cat_info.json'42    print('Saving to', out_path)43    json.dump(out, open(out_path, 'w'))44    45