#!/usr/bin/env python
"""
make_safety_dataset_fix.py
──────────────────────────────────────────────────────────────
• train/val COCO JSON을 읽어 실제 이미지 파일 존재 여부를 검사
• 없는 이미지는 images·annotations 모두 제거
• *_fix.json, missing_*.txt 파일로 결과 저장
"""

import json, os, sys
from pathlib import Path
from typing import Dict, Any
try:
    from tqdm import tqdm
except ImportError:                     # tqdm 없으면 대체
    def tqdm(x, **k): return x

# ── 1. 설정 ──────────────────────────────────────────────────────────
DATA_ROOT   = Path("/")                      # 절대경로 그대로 쓰는 JSON이므로 '/' 고정
TRAIN_JSON  = Path("/data/ai-hub/ppe-dataset/105.공사현장_안전장비_인식_데이터/train_safety_dataset_coco.json")
VAL_JSON    = Path("/data/ai-hub/ppe-dataset/105.공사현장_안전장비_인식_데이터/val_safety_dataset_coco.json")

OUT_TRAIN   = TRAIN_JSON.with_name("train_safety_dataset_coco_fix.json")
OUT_VAL     = VAL_JSON.with_name("val_safety_dataset_coco_fix.json")
MISS_TRAIN  = TRAIN_JSON.with_name("missing_train.txt")
MISS_VAL    = VAL_JSON.with_name("missing_val.txt")

# ── 2. 함수 ──────────────────────────────────────────────────────────
def make_fix(src: Path, dst: Path, miss_txt: Path) -> None:
    coco: Dict[str, Any] = json.load(src.open())
    keep_images = []
    miss_list   = []

    print(f"\n▶︎ scanning {src.name} …")
    for img_meta in tqdm(coco["images"]):
        if os.path.exists(img_meta["file_name"]):
            keep_images.append(img_meta)
        else:
            miss_list.append(img_meta["file_name"])

    # image_id set
    valid_ids = {im["id"] for im in keep_images}

    # filter annotations
    keep_ann = [ann for ann in coco["annotations"] if ann["image_id"] in valid_ids]

    # 요약
    print(f"   total images : {len(coco['images']):,}")
    print(f"   kept images  : {len(keep_images):,}")
    print(f"   missing imgs : {len(miss_list):,}")
    print(f"   kept annot   : {len(keep_ann):,}\n")

    # 새 coco dict
    coco_fix = coco.copy()
    coco_fix["images"] = keep_images
    coco_fix["annotations"] = keep_ann

    # 저장
    json.dump(coco_fix, dst.open("w"), ensure_ascii=False, indent=2)
    print(f"[✓] cleaned JSON  → {dst}")

    if miss_list:
        miss_txt.write_text("\n".join(miss_list))
        print(f"[✓] missing list → {miss_txt}")
    else:
        print("[✓] no missing files")

# ── 3. 실행 ──────────────────────────────────────────────────────────
if __name__ == "__main__":
    for p in (TRAIN_JSON, VAL_JSON):
        if not p.exists():
            sys.exit(f"❌ {p} not found")

    make_fix(TRAIN_JSON, OUT_TRAIN, MISS_TRAIN)
    make_fix(VAL_JSON,   OUT_VAL,   MISS_VAL)