import json
import os

def process_annotations(input_json, output_json, image_dir_1,image_dir_2 default_ext=".jpg"):
    """
    处理注释文件，确保每个图像都包含 file_name 字段，并保存处理后的文件。
    
    Args:
        input_json (str): 输入 JSON 文件路径 (如 LVIS 或 COCO 格式的注释文件)。
        output_json (str): 输出 JSON 文件路径。
        image_dir (str): 图像文件的根目录路径，用于检查 file_name 字段是否存在。
        default_ext (str): 如果 file_name 缺失，使用 image_id 生成的默认文件扩展名（默认为 .jpg）。
    """
    with open(input_json, 'r') as f:
        annotations = json.load(f)
    
    images = annotations.get('images', [])
    modified_images = []

    for image in images:
        # 如果 'file_name' 字段缺失，使用 'id' 字段生成 file_name
        if 'file_name' not in image:
            image_id = image['id']
            file_name = f"{image_id}{default_ext}"
            
            # 检查生成的 file_name 是否存在于指定的图像目录中
            file_path_1 = os.path.join(image_dir_1, file_name)
            file_path_2 = os.path.join(image_dir_2, file_name)
            if os.path.exists(file_path_1) or os.path.exists(file_path_2):
                image['file_name'] = file_name
            else:
                print(f"Warning: Generated file_name '{file_name}' does not exist in '{image_dir}'")
                # 如果生成的文件名在目录中不存在，可以选择跳过该图像或者执行其他操作
                continue  # 跳过该图像
            
        modified_images.append(image)

    # 更新注释字典
    annotations['images'] = modified_images
    
    # 将处理后的注释保存为新的 JSON 文件
    with open(output_json, 'w') as f:
        json.dump(annotations, f)
    
    print(f"Processed annotations saved to {output_json}")

if __name__ == "__main__":
    # 输入参数
    input_json = '/root/Open-GroundingDino-main/lvis_v1_val.json'  # 替换为你 LVIS 或 COCO JSON 文件的路径
    output_json = '/root/Open-GroundingDino-main/lvis_v1_coco_val.json'  # 输出文件路径
    image_dir_1 = '/root/val2017'  # 图像文件所在的目录\
    image_dir_2 = '/root/autodl-tmp/train2017'

    # 调用处理函数
    process_annotations(input_json, output_json, image_dir)
