import json
from natsort import natsorted
import os 

json_path = "/work/yuqian_fu/Ego/data_segswap/egoexo_val_framelevel_newprompt_all_instruction.json"
data_path = "/work/yuqian_fu/Ego/data_segswap"
split_path = "/home/yuqian_fu/Projects/ego-exo4d-relation/correspondence/split.json"
with open(json_path, "r") as fp:
    datas = json.load(fp)
with open(split_path, "r") as fp:
    split = json.load(fp)
takes = split["val"]

datas_final_save = []
for take in takes:
    datas_tmp = []
    gt_path = f"{data_path}/{take}/annotation.json"
    with open(gt_path, 'r') as fp:
        gt = json.load(fp)
    # 创建逆字典
    objs = natsorted(list(gt["masks"].keys()))
    coco_id_to_cont_id = {cont_id + 1: coco_id for cont_id, coco_id in enumerate(objs)}
  

    for data in datas:
        if data["video_name"] == take:
            datas_tmp.append(data)
    
    if len(datas_tmp) == 0:
        print("No data found for take:", take)
        continue
    
    for data in datas_tmp:
        annos = data["first_frame_anns"]
        annos_new = []
        for anno in annos:
            category_id = anno["category_id"]
            obj_name = coco_id_to_cont_id[category_id]
            anno["text"] = obj_name
            annos_new.append(anno)
        data["first_frame_anns"] = annos_new
        del data["instruction"]
        datas_final_save.append(data)





print("datas_final_save:", len(datas_final_save))
save_path = "/work/yuqian_fu/Ego/data_segswap/egoexo_val_framelevel_newprompt_all_objname.json"
with open(save_path, "w") as fp:
    json.dump(datas_final_save, fp)




            





        
# path = "/home/yuqian_fu/Projects/ego-exo4d-relation/correspondence/split.json"
# with open(path, "r") as fp:
#     data = json.load(fp)
# takes = data["test"]
# print("takes:", len(takes))


# takes_new = os.listdir("/work/yuqian_fu/Ego/data_segswap_test_20250508_new888_v2")
# print("takes_new:", len(takes_new))

# # 找出 takes 中有但 takes_new 中没有的元素
# diff_in_takes = set(takes) - set(takes_new)

# # 找出 takes_new 中有但 takes 中没有的元素
# diff_in_takes_new = set(takes_new) - set(takes)

# # 打印结果
# print("takes 中有但 takes_new 中没有的元素:", diff_in_takes)
# print("takes_new 中有但 takes 中没有的元素:", diff_in_takes_new)