import json
import pandas as pd
from pathlib import Path

coco_json = Path('./jsons/test_public.json')
pred_json = Path('./results/out.pkl.json')

img_names = dict()
with coco_json.open() as f:
    for img_ann in json.load(f)['images']:
        img_names[img_ann['id']] = img_ann['file_name']
with pred_json.open() as f:
    preds = json.load(f)

results = []
for pred in preds:
    filename = img_names[pred['image_id']]
    x, y, w, h = pred['bbox']

    results.append({
        'image_filename': filename,
        'label_id': pred['category_id'] + 1,
        'x': x,
        'y': y,
        'w': w,
        'h': h,
        'confidence': pred['score'],
    })

df = pd.DataFrame(results)
df.to_csv('./results/out.csv', index=False,
    columns=['image_filename', 'label_id', 'x', 'y', 'w', 'h', 'confidence'])