# Description: Filter the prediction file to only include the images that are in the ground truth file.
#
# Usage: python filter_test_prediction.py -g <gt_file> -p <prediction_file> -o <output_file>
#   -g, --gt: Path to the ground truth file.
#   -p, --prediction: Path to the prediction file.
#   -o, --output: Path to the output file.
#

import os
import json
import argparse

def main(args):
    with open(args.gt) as f:
        test_gt = json.load(f)

    test_gt = [img['id'] for img in test_gt['images']]

    with open(args.pred) as f:
        test_pred = json.load(f)

    filtered_test_pred = []
    for pred in test_pred:
        if pred['image_id'] in test_gt:
            filtered_test_pred.append(pred)

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


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-g','--gt', type=str, required=True)
    parser.add_argument('-p','--pred', type=str, required=True)
    parser.add_argument('-o','--output', type=str, required=True)
    args = parser.parse_args()

    main(args)