# Description: Make sampled ground truth for evaluation
#
# Usage: python make_sample_gt.py -i <input_file> -o <output_file> -p <prediction_file>
#   -i, --input: Path to the input JSON file.
#   -o, --output: Path to the output JSON file.
#   -p, --prediction: Path to the prediction JSON file.

import json
import argparse

parser = argparse.ArgumentParser(description='Make sample ground truth')
parser.add_argument('-i', '--input', type=str, help='Input file')
parser.add_argument('-o', '--output', type=str, help='Output file')
parser.add_argument('-p', '--prediction', type=str, help='Prediction file')
args = parser.parse_args()


with open(args.input, 'r') as f:
    gt_data = json.load(f)


with open(args.prediction, 'r') as f:
    pred_data = json.load(f)

pred_img = []

for img in pred_data:
    pred_img.append(img['image_id'])

sample_gt = {'images': [], 'annotations': []}
sample_gt['categories'] = gt_data['categories']

for img in gt_data['images']:
    if img['id'] in pred_img:
        sample_gt['images'].append(img)

for ann in gt_data['annotations']:
    if ann['image_id'] in pred_img:
        sample_gt['annotations'].append(ann)

with open(args.output, 'w') as f:
    json.dump(sample_gt, f)