import json
import os
def get_category_dict():
    path = "../POPE/output/coco"
    coco_pope_file_ls = ["coco_pope_adversarial.json","coco_pope_popular.json", "coco_pope_random.json"]
    obj_dict = {}
    for coco_pope_file in coco_pope_file_ls:
        with open(os.path.join(path, coco_pope_file), 'r') as f:
            coco_data = [json.loads(q) for q in f]
        print(len(coco_data), coco_data[0])
        for i in range(0,len(coco_data),6):
            obj_ls = []
            for j in range(6):
                obj = coco_data[i+j]["text"].split(" ")[3:-3]
                # obj set turn to string
                obj = " ".join(obj)
                obj_ls.append(obj)
            obj_ls = set(obj_ls)
            # obj_dict.append({"id": coco_data[i]['question_id'], "image": coco_data[i]['image'], "objects": list(obj_ls)})
            if coco_data[i]['image'] not in obj_dict:
                obj_dict[coco_data[i]['image']] = list(obj_ls)
            else:
                obj_dict[coco_data[i]['image']] += list(obj_ls)
            obj_dict[coco_data[i]['image']] = list(set(obj_dict[coco_data[i]['image']]))
    print(len(obj_dict))
    print(obj_dict['COCO_val2014_000000144305.jpg'])
    with open(os.path.join(path, "minival2014_pope_obj_dict.json"), 'w') as f:
        json.dump(obj_dict, f)
    print(f"save at {os.path.join(path, 'minival2014_pope_obj_dict.json')} successfully!")

import json
import os
import shutil
def get_interest_obj_ls():
    # read file
    file = "../POPE/output/coco/minival2014_pope_obj_dict.json"
    file_copy = file.replace(".json", "_copy.json")
    if not os.path.exists(file_copy):
        shutil.copy(file, file_copy)
    with open(file_copy, 'r') as f:
        interest_obj_ls = json.load(f)
    if len(interest_obj_ls) == 0:
        shutil.copy(file, file_copy)
        with open(file_copy, 'r') as f:
            interest_obj_ls = json.load(f)  
    # pop the first element
    interest_obj = interest_obj_ls.pop(list(interest_obj_ls.keys())[0])
    with open(file_copy, 'w') as f:
        json.dump(interest_obj_ls, f)
    return interest_obj

if __name__ == "__main__":
    # get_category_dict()
    print(get_interest_obj_ls())