# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

import json

with open('./data/lvis_v1/annotations/lvis_v1_train.json') as f:
# with open('lvis_v1/annotations/lvis_v1_val.json') as f:
    data = json.load(f)

print(data.keys())

output_data = [dict(), dict(), dict(), dict()]
for output in output_data:
    output['info'] = data['info']
    output['licenses'] = data['licenses']
    output['categories'] = []
    output['images'] = []
    output['annotations'] = []

images = {}
for img in data['images']:
    images[img['id']] = img

for ann in data['annotations']:
    
    cat_id, img_id = ann['category_id'], ann['image_id']
    category = data['categories'][cat_id-1].copy()
    image = images[img_id].copy()
    annotation = ann.copy()
    
    ins_count = category['instance_count']
    if ins_count < 10:
        output_data[0]['categories'].append(category)
        output_data[0]['images'].append(image)
        output_data[0]['annotations'].append(annotation)
    elif ins_count < 100:
        output_data[1]['categories'].append(category)
        output_data[1]['images'].append(image)
        output_data[1]['annotations'].append(annotation)
    elif ins_count < 1000:
        output_data[2]['categories'].append(category)
        output_data[2]['images'].append(image)
        output_data[2]['annotations'].append(annotation)
    else:
        output_data[3]['categories'].append(category)
        output_data[3]['images'].append(image)
        output_data[3]['annotations'].append(annotation)

for i in range(4):
    L = output_data[i]['categories'] 
    output_data[i]['categories'] = list({v['id']:v for v in L}.values())
    M = output_data[i]['images'] 
    output_data[i]['images'] = list({v['id']:v for v in M}.values())

for i in range(4):
    print(len(output_data[i]['categories']), len(output_data[i]['images']))

for i in range(1,5):
    output_json = './data/lvis_v1/annotations/lvis_v1_train_bin%d.json'%i
    with open(output_json, 'w') as json_file:
        json.dump(output_data[i-1], json_file)





StarcoderdataJupyterScriptsDedupFiltered