import json
import os
# 读取im_rgba.txt中的内容作为bg prompt写入json

# folder_path = '/mnt/workspace/workgroup/sihui.jsh/layer_diff_dataset/try/im'
# folder_path_1 = '/mnt/workspace/workgroup/sihui.jsh/layer_diff_dataset/train/im'

# file_list = os.listdir(folder_path)
# file_list = [os.path.join(folder_path_1,i) for i in file_list if i.endswith('.jpg')]
# file_list.sort()
# print(file_list)
# 读取原始JSON文件
data_folder = '/mnt/workspace/workgroup/sihui.jsh/layer_diff_dataset/train'
with open(os.path.join(data_folder,'im_rgba.json'), 'r') as file:
    data = json.load(file)

captions_dict = {}
with open(os.path.join(data_folder, 'im_rgba.txt'), 'r') as f:
    for line in f.readlines():
        image_name, caption = line.strip().split('\t')
        captions_dict[image_name] = caption
# data_try = []
# # 修改每个字典的"images"和"images_rgba"字段
# for item in data:
#     if item["images"] not in file_list:
#         data_try.append(item)
# print(data_try)
for item in data:
    image_name = item["images"].split('/')[-1]
    # item['prompt'] = class_name + ',' + item['prompt']
    item['bg_prompt'] = captions_dict[image_name]
    # print(item)
    # break
# 写入新的JSON文件
with open('/mnt/workspace/workgroup/sihui.jsh/layer_diff_dataset/train/im_rgba_bg.json', 'w') as file:
    json.dump(data, file, indent=4)

print("New JSON file has been created successfully.")