#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
[목적]
    JSON 파일 내 이미지 정보에서 실제 이미지 파일이 존재하지 않는 항목을 제거하고,
    해당 이미지 id를 참조하는 annotation도 함께 삭제하여 새 JSON 파일로 저장하는 스크립트입니다.

[사용 안내]
    1. 아래 코드를 filter_missing_images.py 라는 파일명으로 저장합니다.
    2. 터미널에서 다음과 같이 실행합니다:
           python filter_missing_images.py --input_json /절대경로/원본.json --output_json /절대경로/필터된.json
    3. 스크립트는 JSON 파일의 "images" 항목에서 각 이미지의 file_name (또는 img_path) 경로를 확인하여,
       실제 파일이 존재하지 않으면 해당 항목과 관련 annotation을 제거합니다.
    4. 최종 결과물은 지정한 output_json 파일로 저장됩니다.

[코드 설명]
    - argparse를 사용하여 입력 및 출력 JSON 파일 경로를 인자로 받습니다.
    - JSON 파일을 읽고, "images" 리스트를 순회하며 os.path.exists()로 파일 존재 여부를 확인합니다.
    - 존재하는 이미지의 id만 valid_image_ids 리스트에 저장하고, 존재하지 않는 항목은 missing_images 리스트에 모읍니다.
    - annotation도 valid_image_ids에 해당하는 것들만 필터링합니다.
    - 최종적으로 filtered된 JSON 데이터를 지정한 경로에 저장하고, 요약 정보를 출력합니다.
"""

import os
import json
import argparse

def filter_missing_images(input_json_path, output_json_path):
    # JSON 파일 로드
    with open(input_json_path, 'r', encoding='utf-8') as f:
        data = json.load(f)
    
    images = data.get("images", [])
    annotations = data.get("annotations", [])
    
    valid_image_ids = []      # 존재하는 이미지의 id 목록
    filtered_images = []      # 존재하는 이미지 정보
    missing_images = []       # 존재하지 않는 이미지 정보
    
    # 각 이미지의 파일 존재 여부 확인
    for img in images:
        # file_name 또는 img_path (둘 중 하나가 사용될 수 있음)
        file_path = img.get("file_name", "") or img.get("img_path", "")
        if os.path.exists(file_path):
            filtered_images.append(img)
            valid_image_ids.append(img["id"])
        else:
            missing_images.append(img)
    
    # valid_image_ids에 포함되지 않은 annotation 제거
    filtered_annotations = [ann for ann in annotations if ann["image_id"] in valid_image_ids]
    
    # 데이터 업데이트
    data["images"] = filtered_images
    data["annotations"] = filtered_annotations
    
    # 요약 정보 출력
    print("총 이미지 개수:", len(images))
    print("유효한 이미지 개수:", len(filtered_images))
    print("제거된 이미지 개수:", len(missing_images))
    
    if missing_images:
        print("\n존재하지 않는 이미지 파일 목록:")
        for img in missing_images:
            print("  -", img.get("file_name", img.get("img_path", "Unknown")))
    
    # 필터링된 결과를 새 JSON 파일로 저장
    with open(output_json_path, 'w', encoding='utf-8') as f:
        json.dump(data, f, ensure_ascii=False, indent=4)
    
    print(f"\n필터링된 JSON 파일이 저장되었습니다: {output_json_path}")

def main():
    parser = argparse.ArgumentParser(description="JSON 파일에서 존재하지 않는 이미지 항목 제거 스크립트")
    parser.add_argument("--input_json", type=str, required=True, help="입력 JSON 파일의 절대 경로")
    parser.add_argument("--output_json", type=str, required=True, help="출력 JSON 파일의 절대 경로 (새 파일로 저장됨)")
    args = parser.parse_args()
    filter_missing_images(args.input_json, args.output_json)

if __name__ == "__main__":
    main()