#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
remove_broken_from_coco.py
──────────────────────────────────────────────────────────────
  • COCO JSON(대용량)에서 깨진 이미지 목록(broken.txt)에 해당하는
    image·annotation 모두 삭제
  • *_fix.json 으로 저장
  • 삭제/유지 통계 출력

사용 방법
────────
1) 아래 변수만 경로에 맞게 수정
2) python remove_broken_from_coco.py
   ├─ xx_fix.json  (새 JSON)
   └─ console 에 통계 출력
"""

import json
from pathlib import Path
from tqdm import tqdm

# ─── 1. 경로 설정 ────────────────────────────────────────────────────
JSON_PATH   = Path("/data/ai-hub/ppe-dataset/105.공사현장_안전장비_인식_데이터/train_safety_dataset_coco_fix_broken_fix.json")
BROKEN_LIST = JSON_PATH.with_suffix(".broken.txt")  # ← 앞서 생성된 파일
OUT_JSON    = JSON_PATH.with_name("train_safety_dataset_coco_fix_broken_fix.json")

# ─── 2. 깨진 파일 목록 읽기 → set -------------------------------------------------
broken_set = {line.strip() for line in BROKEN_LIST.read_text().splitlines() if line.strip()}
print(f"🔹 broken list  : {len(broken_set):,} files")

# ─── 3. COCO JSON 로드 -----------------------------------------------------------
print(f"🔹 loading COCO : {JSON_PATH.name}")
coco = json.load(JSON_PATH.open())

# ─── 4. images 필터링 ------------------------------------------------------------
keep_images = []
broken_ids  = set()

for img in tqdm(coco["images"], desc="filter images"):
    if img["file_name"] in broken_set:
        broken_ids.add(img["id"])
    else:
        keep_images.append(img)

print(f"\n   kept images  : {len(keep_images):,}")
print(f"   removed imgs : {len(broken_ids):,}")

# ─── 5. annotations 필터링 --------------------------------------------------------
keep_ann = [ann for ann in coco["annotations"] if ann["image_id"] not in broken_ids]

print(f"   kept annots  : {len(keep_ann):,}")
print(f"   removed ann  : {len(coco['annotations']) - len(keep_ann):,}")

# ─── 6. 새 JSON 저장 --------------------------------------------------------------
coco_fix = coco.copy()
coco_fix["images"]       = keep_images
coco_fix["annotations"]  = keep_ann

json.dump(coco_fix, OUT_JSON.open("w"), ensure_ascii=False, indent=2)
print(f"\n[✓] saved → {OUT_JSON.resolve()}")