{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %% =====================================================================\n",
    "# [Cell 1] 초기 설정 – import & 전역 변수 (공사현장 안전장비용)\n",
    "# =====================================================================\n",
    "import os, json, random, cv2\n",
    "from pathlib import Path\n",
    "from collections import defaultdict\n",
    "from tqdm.auto import tqdm\n",
    "import numpy as np\n",
    "\n",
    "# ---------------- 사용자 설정 ----------------\n",
    "DATA_ROOT = Path(\"/data/ai-hub/ppe-dataset/105.공사현장_안전장비_인식_데이터/01.데이터\")\n",
    "\n",
    "# (1) 어노테이션 JSON 루트 – Training / Validation 두 곳\n",
    "ANNOT_DIRS = [\n",
    "    DATA_ROOT / \"1.Training/라벨링데이터_241008_add\",\n",
    "    # DATA_ROOT / \"2.Validation/라벨링데이터(zip)_241008_add\",\n",
    "]\n",
    "\n",
    "# (2) 이미지(원천데이터) 루트\n",
    "IMG_DIRS = [\n",
    "    DATA_ROOT / \"1.Training/원천데이터_210818_add\",\n",
    "    # DATA_ROOT / \"2.Validation/원천데이터_210818_add\",\n",
    "]\n",
    "\n",
    "# (3) 클래스 매핑 – 45 종 (id는 1부터 45까지)\n",
    "CLASS_MAPPING = {\n",
    "    f\"{i:02d}\": {\"id\": i, \"name\": name} for i, name in enumerate([\n",
    "        # 01–08 안전보호구\n",
    "        \"belt_on\", \"belt_off\", \"hook_on\", \"hook_off\",\n",
    "        \"shoe_on\", \"shoe_off\", \"helmet_on\", \"helmet_off\",\n",
    "        # 09–19 중장비\n",
    "        \"excavator\", \"payloader\", \"forklift\", \"dump_truck\", \"concrete_mixer\",\n",
    "        \"pump_car\", \"pile_driver\", \"truck\", \"aerial_platform\",\n",
    "        \"tower_crane\", \"sky_lift\",\n",
    "        # 20–21 구조물\n",
    "        \"gang_form\", \"al_form\",\n",
    "        # 22–23 사다리\n",
    "        \"ladder_A\", \"ladder_worktable\",\n",
    "        # 24–29 설치물\n",
    "        \"switch_box\", \"cover_opening\", \"hazard_storage\",\n",
    "        \"elev_falling_protection\", \"hoist\", \"jack_support\",\n",
    "        # 30–31 비계\n",
    "        \"pipe_scaffold\", \"system_scaffold\",\n",
    "        # 32–45 자재·공구\n",
    "        \"cement_brick\", \"hammer\", \"drill\", \"remital\", \"brick_finish\",\n",
    "        \"mixer\", \"H_beam\", \"cutting_machine\", \"vibrator\",\n",
    "        \"fire_extinguisher\", \"welder\", \"grinder\",\n",
    "        \"hand_cart\", \"spark_protection\",\n",
    "    ], start=1)\n",
    "}\n",
    "\n",
    "# (4) 출력 경로\n",
    "OUTPUT_JSON = DATA_ROOT.parent / \"train_safety_dataset_coco.json\"\n",
    "\n",
    "RANDOM_SEED = 42\n",
    "print(\"✅ 설정 완료\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %% =====================================================================\n",
    "# [셀 2] 이미지 인덱스 + JSON 폴더 검증\n",
    "# =====================================================================\n",
    "from glob import glob\n",
    "\n",
    "# 2-A) 어노테이션 폴더마다 JSON 개수 확인\n",
    "for d in ANNOT_DIRS:\n",
    "    n = len(list(Path(d).rglob(\"*.json\")))\n",
    "    if n == 0:\n",
    "        print(f\"⚠️  경고: {d} 에서 JSON을 찾지 못했습니다 (경로 오타 가능)\")\n",
    "    else:\n",
    "        print(f\"✅ {d} → JSON {n:,}개\")\n",
    "\n",
    "# 2-B) 이미지 인덱싱\n",
    "def build_img_idx(dirs):\n",
    "    by_name, by_stem = {}, {}\n",
    "    for root in dirs:\n",
    "        print(\"📂 이미지 탐색:\", root)\n",
    "        for ext in (\"*.jpg\",\"*.jpeg\",\"*.png\",\"*.JPG\",\"*.PNG\"):\n",
    "            for fp in glob(str(root / \"**\" / ext), recursive=True):\n",
    "                p = Path(fp)\n",
    "                by_name.setdefault(p.name, fp)\n",
    "                by_stem.setdefault(p.stem, fp)\n",
    "        print(f\"   ↳ 누적 {len(by_name):,} files\")\n",
    "    print(f\"✅ 인덱싱 완료 – 총 {len(by_name):,}\")\n",
    "    return by_name, by_stem\n",
    "\n",
    "IMG_BY_NAME, IMG_BY_STEM = build_img_idx(IMG_DIRS)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %% =====================================================================\n",
    "# [셀 3] JSON 로더 – 어떤 포맷이 와도 꿀꺽!\n",
    "# =====================================================================\n",
    "import json, cv2, numpy as np\n",
    "from pathlib import Path\n",
    "\n",
    "def safe_load(p):\n",
    "    try:\n",
    "        return json.load(open(p, encoding=\"utf-8-sig\"))\n",
    "    except Exception as e:\n",
    "        print(\"⚠️ JSON 오류:\", p, e); return None\n",
    "\n",
    "def to_bbox(box_like):\n",
    "    \"\"\"다양한 bbox 표현을 [x,y,w,h] 로 통일\"\"\"\n",
    "    if not box_like: return None\n",
    "    # 리스트 [x1,y1,x2,y2]\n",
    "    if isinstance(box_like, list) and len(box_like) == 4:\n",
    "        x1,y1,x2,y2 = map(float, box_like)\n",
    "        return [x1,y1,x2-x1,y2-y1]\n",
    "    # dict 형태\n",
    "    if isinstance(box_like, dict):\n",
    "        if {\"x\",\"y\",\"w\",\"h\"} <= box_like.keys():\n",
    "            return [box_like[\"x\"], box_like[\"y\"], box_like[\"w\"], box_like[\"h\"]]\n",
    "        if {\"xmin\",\"ymin\",\"xmax\",\"ymax\"} <= box_like.keys():\n",
    "            x1,y1,x2,y2 = [box_like[k] for k in (\"xmin\",\"ymin\",\"xmax\",\"ymax\")]\n",
    "            return [x1,y1,x2-x1,y2-y1]\n",
    "    return None\n",
    "\n",
    "def poly2bbox(poly):\n",
    "    pts = np.asarray(poly, dtype=np.float32).reshape(-1,2)\n",
    "    x0,y0 = pts.min(0); x1,y1 = pts.max(0)\n",
    "    return [float(x0), float(y0), float(x1-x0), float(y1-y0)]\n",
    "\n",
    "def find_img_any(name_or_path):\n",
    "    \"\"\"파일명·stem·상대경로 어떤 것이 와도 IMG 인덱스에서 찾아줌\"\"\"\n",
    "    p = Path(name_or_path)\n",
    "    return (IMG_BY_NAME.get(p.name) or\n",
    "            IMG_BY_STEM.get(p.stem))\n",
    "\n",
    "def load_safe_json(js_path: Path):\n",
    "    d = safe_load(js_path)\n",
    "    if not d: return None, None\n",
    "\n",
    "    # ---------- 파일명 ----------\n",
    "    if \"image\" in d:\n",
    "        fname = d[\"image\"].get(\"filename\") or d[\"image\"].get(\"file_name\")\n",
    "    else:\n",
    "        # image 블록이 없는 경우 → JSON 이름으로 추정\n",
    "        fname = js_path.stem + \".jpg\"\n",
    "    img_path = find_img_any(fname)\n",
    "    if img_path is None:\n",
    "        return None, None   # 인덱스에도 없으면 skip\n",
    "\n",
    "    # ---------- 해상도 ----------\n",
    "    if \"image\" in d:\n",
    "        meta = d[\"image\"]\n",
    "        w = meta.get(\"width\")\n",
    "        h = meta.get(\"height\")\n",
    "        if w is None and meta.get(\"resolution\"):\n",
    "            w, h = meta[\"resolution\"]\n",
    "    else:\n",
    "        w = h = None\n",
    "    if w is None or h is None:\n",
    "        img = cv2.imread(img_path)\n",
    "        if img is None: return None, None\n",
    "        h, w = img.shape[:2]\n",
    "\n",
    "    # ---------- 어노테이션 ----------\n",
    "    ann_raw = d.get(\"annotations\") or d.get(\"objects\") or d.get(\"labels\") or []\n",
    "    parsed = []\n",
    "    for a in ann_raw:\n",
    "        cls = a.get(\"class\") or a.get(\"label\") or a.get(\"category_id\")\n",
    "        if isinstance(cls, int): cls = f\"{cls:02d}\"\n",
    "        if cls not in CLASS_MAPPING: continue\n",
    "\n",
    "        bbox = to_bbox(a.get(\"bbox\") or a.get(\"box\")) \\\n",
    "               or (poly2bbox(a[\"polygon\"]) if \"polygon\" in a else None)\n",
    "        if bbox is None: continue\n",
    "\n",
    "        parsed.append({\"class_code\": cls, \"bbox\": bbox})\n",
    "\n",
    "    if not parsed: return None, None\n",
    "    return {\"filename\": fname, \"width\": w, \"height\": h}, parsed\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %% =====================================================================\n",
    "# [셀 4] COCO 변환 – 매칭 실패는 Warning 만\n",
    "# =====================================================================\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "def convert_to_coco(annot_dirs):\n",
    "    images, annotations = [], []\n",
    "    img_id = ann_id = 0\n",
    "\n",
    "    json_files = []\n",
    "    for d in annot_dirs:\n",
    "        json_files += list(Path(d).rglob(\"*.json\"))\n",
    "    random.seed(RANDOM_SEED); random.shuffle(json_files)\n",
    "\n",
    "    pbar = tqdm(json_files, desc=\"JSON → COCO\", ncols=110)\n",
    "    for js in pbar:\n",
    "        meta, anns = load_safe_json(js)\n",
    "        if meta is None:          # 파싱 실패 or 이미지 없음\n",
    "            continue\n",
    "\n",
    "        img_id += 1\n",
    "        images.append({\n",
    "            \"id\"       : img_id,\n",
    "            \"file_name\": find_img_any(meta[\"filename\"]),\n",
    "            \"width\"    : meta[\"width\"],\n",
    "            \"height\"   : meta[\"height\"],\n",
    "        })\n",
    "\n",
    "        for a in anns:\n",
    "            ann_id += 1\n",
    "            cid  = CLASS_MAPPING[a[\"class_code\"]][\"id\"]\n",
    "            bbox = a[\"bbox\"]\n",
    "            annotations.append({\n",
    "                \"id\"         : ann_id,\n",
    "                \"image_id\"   : img_id,\n",
    "                \"category_id\": cid,\n",
    "                \"bbox\"       : bbox,\n",
    "                \"area\"       : bbox[2]*bbox[3],\n",
    "                \"iscrowd\"    : 0,\n",
    "            })\n",
    "\n",
    "    print(f\"✅ 변환 완료 – 이미지 {len(images):,}, bbox {len(annotations):,}\")\n",
    "    return {\n",
    "        \"info\"       : {\"description\": \"AI-Hub Construction-Safety Dataset (COCO)\"},\n",
    "        \"licenses\"   : [],\n",
    "        \"categories\" : list(CLASS_MAPPING.values()),\n",
    "        \"images\"     : images,\n",
    "        \"annotations\": annotations,\n",
    "    }\n",
    "\n",
    "COCO_DATA = convert_to_coco(ANNOT_DIRS)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %% =====================================================================\n",
    "# [Cell 5] COCO JSON 저장\n",
    "# =====================================================================\n",
    "OUTPUT_JSON.parent.mkdir(parents=True, exist_ok=True)\n",
    "with OUTPUT_JSON.open(\"w\", encoding=\"utf-8\") as f:\n",
    "    json.dump(COCO_DATA, f, ensure_ascii=False, indent=2)\n",
    "print(f\"📦 저장 완료 → {OUTPUT_JSON}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %% =====================================================================\n",
    "# [Cell 6] (선택) 랜덤 5장 시각화 & 검증\n",
    "# =====================================================================\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "DEBUG_VIS_NUM = 8\n",
    "sample_imgs = random.sample(COCO_DATA[\"images\"], min(DEBUG_VIS_NUM, len(COCO_DATA[\"images\"])))\n",
    "anns_by_img = defaultdict(list)\n",
    "for a in COCO_DATA[\"annotations\"]:\n",
    "    anns_by_img[a[\"image_id\"]].append(a)\n",
    "\n",
    "for img_meta in sample_imgs:\n",
    "    img = cv2.imread(img_meta[\"file_name\"])\n",
    "    if img is None:\n",
    "        print(\"⚠️ load fail\", img_meta[\"file_name\"]); continue\n",
    "    for a in anns_by_img[img_meta[\"id\"]]:\n",
    "        cid = a[\"category_id\"]; x,y,w,h = map(int, a[\"bbox\"])\n",
    "        color = ((cid*37)%256, (cid*17)%256, (cid*13)%256)\n",
    "        cv2.rectangle(img,(x,y),(x+w,y+h),color,2)\n",
    "        cv2.putText(img, str(cid), (x, y-4), cv2.FONT_HERSHEY_SIMPLEX, .6, color, 2)\n",
    "\n",
    "    # plt.figure(figsize=(6,6)); plt.axis(\"off\")\n",
    "    plt.figure(figsize=(24,24), dpi=200); plt.axis(\"off\")\n",
    "    plt.title(Path(img_meta[\"file_name\"]).name)\n",
    "    plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB)); plt.show()\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "250520_mmyolo",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.20"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}