{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[Cell 1] Imports 및 글로벌 설정\n",
    "DATASET_PATHS: 합치고 싶은 YOLO11 데이터셋 폴더 경로들을 리스트에 추가\n",
    "\n",
    "OUTPUT_BASE: 합쳐진 train/val을 저장할 최상위 경로\n",
    "\n",
    "TRAIN_RATIO 등 분할 비율, 시드, 디버깅 옵션 설정\n",
    "\n",
    "[Cell 2] 유틸리티 함수 정의\n",
    "ensure_dir, read_label_file, collect_all_samples 등\n",
    "\n",
    "각 데이터셋을 합쳐서 샘플 수집, 클래스 카운트, 고유명 생성 로직 포함\n",
    "\n",
    "[Cell 3] Train/Val 분할 및 복사\n",
    "split_and_copy()로 랜덤 분할 → val에 누락 클래스 보강 → 파일 복사\n",
    "\n",
    "tqdm 진행바, 예외 상황만 로깅, 마지막에 요약 통계 출력\n",
    "\n",
    "[Cell 4] Debug Visualization (옵션)\n",
    "DEBUG_VIS=True일 때, train/val에서 각각 DEBUG_VIS_COUNT장씩 골라\n",
    "\n",
    "draw_annotations()로 bbox+클래스명 시각화 후 debug_vis 폴더에 저장\n",
    "\n",
    "이제 지정한 여러 데이터셋을 한 번에 합쳐 train/val로 분류하고,\n",
    "파일명 중복 방지 및 오류 로깅, 디버깅용 시각화까지 자동으로 처리하실 수 있습니다.\n",
    "필요한 경로만 수정하시면 되며, 추가 문의 주시면 바로 도와드리겠습니다!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# [Cell 1] Imports 및 글로벌 설정\n",
    "import os\n",
    "import random\n",
    "import shutil\n",
    "from collections import defaultdict\n",
    "import cv2\n",
    "\n",
    "# ---------------------------\n",
    "# (1) 데이터셋 구성: 경로 및 로컬 클래스 매핑\n",
    "# ---------------------------\n",
    "DATASET_CONFIGS = [\n",
    "    # {\n",
    "    #     'path': \"/ltb/media/ltb/90887173887158A4/Users/ltb/박스_데이터셋/박스_yolo11/it1_1819_4class_whhb\",\n",
    "    #     'class_names': {0: 'worker', 1: 'helmet', 2: 'head', 3: 'background'}\n",
    "    # },\n",
    "\n",
    "    # {\n",
    "    #     'path': \"/ltb/media/ltb/90887173887158A4/Users/ltb/박스_데이터셋/박스_yolo11/it1_1819_4class_whhb_albu_aug\",\n",
    "    #     'class_names': {0: 'worker', 1: 'helmet', 2: 'head', 3: 'background'}\n",
    "    # },\n",
    "\n",
    "    {\n",
    "        'path': \"/media/labelstudio/files/박스_yolo11/dl_all_combined\",\n",
    "        'class_names': {0: 'worker', \n",
    "                        1: 'signalman', 2: 'helmet', 3: 'harness',\n",
    "                        4: 'mixer_truck',\n",
    "                        5: 'excavator',\n",
    "                        }\n",
    "    },\n",
    "\n",
    "    # 추가 데이터셋은 아래와 같이 설정\n",
    "    # {\n",
    "    #   'path': '/path/to/another',\n",
    "    #   'class_names': {0:'cat',1:'dog'}\n",
    "    # },\n",
    "]\n",
    "\n",
    "# ---------------------------\n",
    "# (2) 분할 비율 및 랜덤 시드\n",
    "# ---------------------------\n",
    "TRAIN_RATIO = 0.8\n",
    "RANDOM_SEED = 42\n",
    "\n",
    "# ---------------------------\n",
    "# (3) 출력 폴더 경로\n",
    "# ---------------------------\n",
    "OUTPUT_BASE     = \"/media/labelstudio/files/박스_yolo11/Box_train_val/250515_dl_all_combined\"\n",
    "IMAGE_TRAIN_DIR = os.path.join(OUTPUT_BASE, 'images', 'train')\n",
    "IMAGE_VAL_DIR   = os.path.join(OUTPUT_BASE, 'images', 'val')\n",
    "LABEL_TRAIN_DIR = os.path.join(OUTPUT_BASE, 'labels', 'train')\n",
    "LABEL_VAL_DIR   = os.path.join(OUTPUT_BASE, 'labels', 'val')\n",
    "MAPPING_TXT_PATH= os.path.join(OUTPUT_BASE, 'class_mapping.txt')\n",
    "\n",
    "# ---------------------------\n",
    "# (4) 디버깅 시각화 옵션\n",
    "# ---------------------------\n",
    "DEBUG_VIS       = True\n",
    "DEBUG_VIS_COUNT = 10\n",
    "DEBUG_VIS_DIR   = os.path.join(OUTPUT_BASE, 'debug_vis')\n",
    "\n",
    "# ---------------------------\n",
    "# (5) 컬러 팔레트 및 텍스트 오프셋\n",
    "# ---------------------------\n",
    "COLOR_PALETTE = [(0,0,255),(0,255,0),(255,0,0),(0,255,255),(255,0,255),(255,255,0)]\n",
    "TEXT_OFFSETS  = {}  # 클래스별 (x,y) 오프셋\n",
    "\n",
    "# ---------------------------\n",
    "# (6) 전역 클래스 통합 매핑 생성\n",
    "# ---------------------------\n",
    "all_names = []\n",
    "for cfg in DATASET_CONFIGS:\n",
    "    for name in cfg['class_names'].values():\n",
    "        if name not in all_names:\n",
    "            all_names.append(name)\n",
    "GLOBAL_CLASS_NAMES = {i: nm for i, nm in enumerate(all_names)}\n",
    "GLOBAL_NAME_TO_INDEX = {nm: i for i, nm in GLOBAL_CLASS_NAMES.items()}\n",
    "\n",
    "os.makedirs(OUTPUT_BASE, exist_ok=True)\n",
    "with open(MAPPING_TXT_PATH, 'w', encoding='utf-8') as mf:\n",
    "    for idx, nm in GLOBAL_CLASS_NAMES.items(): mf.write(f\"{idx}\\t{nm}\\n\")\n",
    "print(\"✅ 글로벌 클래스 매핑:\")\n",
    "for idx, nm in GLOBAL_CLASS_NAMES.items(): print(f\"  {idx}: {nm}\")\n",
    "print(f\"매핑 파일 저장: {MAPPING_TXT_PATH}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# [Cell 2] 유틸리티 함수 정의\n",
    "\n",
    "def ensure_dir(d):\n",
    "    if not os.path.exists(d): os.makedirs(d, exist_ok=True)\n",
    "\n",
    "\n",
    "def read_yolo_labels(path):\n",
    "    ann = []\n",
    "    try:\n",
    "        with open(path, 'r', encoding='utf-8') as f:\n",
    "            for line in f:\n",
    "                t = line.strip().split()\n",
    "                if len(t) != 5: continue\n",
    "                ann.append((int(t[0]), float(t[1]), float(t[2]), float(t[3]), float(t[4])))\n",
    "    except Exception as e:\n",
    "        print(f\"❗ 라벨 읽기 실패: {path} / {e}\")\n",
    "    return ann\n",
    "\n",
    "\n",
    "def convert_to_global(annots, local_map):\n",
    "    ga = []\n",
    "    counts = defaultdict(int)\n",
    "    for lc, xc, yc, w, h in annots:\n",
    "        name = local_map.get(lc)\n",
    "        if name is None: continue\n",
    "        gi = GLOBAL_NAME_TO_INDEX[name]\n",
    "        ga.append((gi, xc, yc, w, h))\n",
    "        counts[gi] += 1\n",
    "    return ga, counts\n",
    "\n",
    "\n",
    "def collect_all_samples(configs):\n",
    "    samples = []\n",
    "    overall = defaultdict(int)\n",
    "    for cfg in configs:\n",
    "        ds = cfg['path']; ds_name = os.path.basename(ds)\n",
    "        img_dir = os.path.join(ds, 'images')\n",
    "        lbl_dir = os.path.join(ds, 'labels')\n",
    "        for lbl in os.listdir(lbl_dir):\n",
    "            if not lbl.lower().endswith('.txt'): continue\n",
    "            base, _ = os.path.splitext(lbl)\n",
    "            # 이미지 찾기\n",
    "            found = None\n",
    "            for ext in ['.jpg', '.png', '.jpeg']:\n",
    "                p = os.path.join(img_dir, base + ext)\n",
    "                if os.path.exists(p): found = p; break\n",
    "            if not found:\n",
    "                print(f\"❗ 이미지 없음: {ds_name}/{lbl}\")\n",
    "                continue\n",
    "            ann = read_yolo_labels(os.path.join(lbl_dir, lbl))\n",
    "            ga, gcounts = convert_to_global(ann, cfg['class_names'])\n",
    "            if not ga: continue\n",
    "            unique_base = f\"{ds_name}_{base}\"\n",
    "            samples.append({\n",
    "                'base': unique_base,\n",
    "                'image_path': found,\n",
    "                'label_path': os.path.join(lbl_dir, lbl),\n",
    "                'annotations': ga,\n",
    "                'class_counts': gcounts\n",
    "            })\n",
    "            for gi, cnt in gcounts.items(): overall[gi] += cnt\n",
    "    return samples, dict(overall)\n",
    "\n",
    "\n",
    "def select_test_samples(samples, n):\n",
    "    allc = set()\n",
    "    for s in samples: allc.update(s['class_counts'].keys())\n",
    "    if n < len(allc):\n",
    "        print(f\"▶ test 수({n}) < 클래스 수({len(allc)}), {len(allc)}로 조정\")\n",
    "        n = len(allc)\n",
    "    sel = []; covered = set()\n",
    "    for s in samples:\n",
    "        new = set(s['class_counts'].keys()) - covered\n",
    "        if new:\n",
    "            sel.append(s);\n",
    "            covered |= set(s['class_counts'].keys())\n",
    "        if covered == allc: break\n",
    "    rem = [s for s in samples if s not in sel]\n",
    "    random.shuffle(rem)\n",
    "    while len(sel) < n and rem:\n",
    "        sel.append(rem.pop())\n",
    "    return sel\n",
    "\n",
    "\n",
    "def draw_annotations(img, annots, names, palette, offsets):\n",
    "    h, w = img.shape[:2]\n",
    "    for c, xc, yc, ww, hh in annots:\n",
    "        color = palette[int(c) % len(palette)]\n",
    "        x1 = int((xc - ww/2) * w); y1 = int((yc - hh/2) * h)\n",
    "        x2 = int((xc + ww/2) * w); y2 = int((yc + hh/2) * h)\n",
    "        cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)\n",
    "        off = offsets.get(int(c), (0, -10))\n",
    "        pos = (x1 + off[0], y1 + off[1])\n",
    "        cv2.putText(img, names.get(int(c), str(c)), pos,\n",
    "                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2)\n",
    "    return img\n",
    "\n",
    "print(\"✅ Cell2: 유틸 함수 정의 완료\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# [Cell 3] Train/Val 분할 및 복사 + 클래스별 통계\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "def split_and_copy(samples, ratio, seed):\n",
    "    random.seed(seed)\n",
    "    random.shuffle(samples)\n",
    "    N = len(samples); n_tr = int(N * ratio)\n",
    "    train = samples[:n_tr]; val = samples[n_tr:]\n",
    "    # val 누락 클래스 보강\n",
    "    allc = set(c for s in samples for c in s['class_counts'])\n",
    "    valc = set(c for s in val for c in s['class_counts'])\n",
    "    for cls in allc - valc:\n",
    "        for s in train:\n",
    "            if cls in s['class_counts']:\n",
    "                val.append(s); train.remove(s)\n",
    "                break\n",
    "    # 출력 폴더 생성\n",
    "    for d in [IMAGE_TRAIN_DIR, IMAGE_VAL_DIR, LABEL_TRAIN_DIR, LABEL_VAL_DIR]:\n",
    "        ensure_dir(d)\n",
    "    stats = {'train':0, 'val':0, 'fail':0}\n",
    "    for group, d_img, d_lbl, tag in [\n",
    "        (train, IMAGE_TRAIN_DIR, LABEL_TRAIN_DIR, 'train'),\n",
    "        (val,   IMAGE_VAL_DIR,   LABEL_VAL_DIR,   'val')\n",
    "    ]:\n",
    "        for s in tqdm(group, desc=f'Copy {tag}', ncols=80):\n",
    "            try:\n",
    "                ext = os.path.splitext(s['image_path'])[1]\n",
    "                img_dst = os.path.join(d_img, s['base'] + ext)\n",
    "                lbl_dst = os.path.join(d_lbl, s['base'] + '.txt')\n",
    "                if not os.path.exists(img_dst): shutil.copy2(s['image_path'], img_dst)\n",
    "                if not os.path.exists(lbl_dst):\n",
    "                    with open(lbl_dst, 'w', encoding='utf-8') as fw:\n",
    "                        for gi, xc, yc, ww, hh in s['annotations']:\n",
    "                            fw.write(f\"{gi} {xc:.6f} {yc:.6f} {ww:.6f} {hh:.6f}\\n\")\n",
    "                stats[tag] += 1\n",
    "            except Exception as e:\n",
    "                print(f\"❗ 복사 실패 ({s['base']}): {e}\")\n",
    "                stats['fail'] += 1\n",
    "    return stats, train, val\n",
    "\n",
    "# 실행 및 요약\n",
    "samples, overall_counts = collect_all_samples(DATASET_CONFIGS)\n",
    "print(f\"총 샘플 수: {len(samples)}\")\n",
    "stats, train, val = split_and_copy(samples, TRAIN_RATIO, RANDOM_SEED)\n",
    "print(\"\\n=== Split Summary ===\")\n",
    "print(f\"Train 복사: {stats['train']} | Val 복사: {stats['val']} | Fail: {stats['fail']}\")\n",
    "print(f\"train 파일 수: {len(train)} | val 파일 수: {len(val)}\")\n",
    "\n",
    "def aggregate_counts(samps):\n",
    "    cnt = defaultdict(int)\n",
    "    for s in samps:\n",
    "        for cls, c in s['class_counts'].items(): cnt[int(cls)] += c\n",
    "    return cnt\n",
    "\n",
    "full_cnt = overall_counts\n",
    "train_cnt = aggregate_counts(train)\n",
    "val_cnt = aggregate_counts(val)\n",
    "\n",
    "print(\"\\n전체 샘플 클래스별 annotation 개수:\")\n",
    "for cls in sorted(GLOBAL_CLASS_NAMES): print(f\"  {cls}({GLOBAL_CLASS_NAMES[cls]}): {full_cnt.get(cls,0)}\")\n",
    "print(\"\\nTrain 셋 클래스별 annotation 개수:\")\n",
    "for cls in sorted(GLOBAL_CLASS_NAMES): print(f\"  {cls}({GLOBAL_CLASS_NAMES[cls]}): {train_cnt.get(cls,0)}\")\n",
    "print(\"\\nVal 셋 클래스별 annotation 개수:\")\n",
    "for cls in sorted(GLOBAL_CLASS_NAMES): print(f\"  {cls}({GLOBAL_CLASS_NAMES[cls]}): {val_cnt.get(cls,0)}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# [Cell 4] Debug Visualization\n",
    "if DEBUG_VIS:\n",
    "    ensure_dir(DEBUG_VIS_DIR)\n",
    "    dbg_samples = []\n",
    "    for s in train:\n",
    "        ext = os.path.splitext(s['image_path'])[1]\n",
    "        dbg_samples.append({\n",
    "            'subset': 'train', 'base': s['base'],\n",
    "            'image': os.path.join(IMAGE_TRAIN_DIR, s['base']+ext),\n",
    "            'label': os.path.join(LABEL_TRAIN_DIR, s['base']+'.txt'),\n",
    "            'class_counts': s['class_counts'],\n",
    "            'annotations': s['annotations']\n",
    "        })\n",
    "    for s in val:\n",
    "        ext = os.path.splitext(s['image_path'])[1]\n",
    "        dbg_samples.append({\n",
    "            'subset': 'val', 'base': s['base'],\n",
    "            'image': os.path.join(IMAGE_VAL_DIR, s['base']+ext),\n",
    "            'label': os.path.join(LABEL_VAL_DIR, s['base']+'.txt'),\n",
    "            'class_counts': s['class_counts'],\n",
    "            'annotations': s['annotations']\n",
    "        })\n",
    "    sel = select_test_samples(dbg_samples, DEBUG_VIS_COUNT)\n",
    "    for s in sel:\n",
    "        img = cv2.imread(s['image'])\n",
    "        if img is None:\n",
    "            print(f\"❗ Debug 시각화 이미지 로드 실패: {s['image']}\")\n",
    "            continue\n",
    "        vis = draw_annotations(img.copy(), s['annotations'], GLOBAL_CLASS_NAMES, COLOR_PALETTE, TEXT_OFFSETS)\n",
    "        outp = os.path.join(DEBUG_VIS_DIR, f\"vis_{s['subset']}_{s['base']}\" + os.path.splitext(s['image'])[1])\n",
    "        cv2.imwrite(outp, vis)\n",
    "    print(f\"✅ Debug 시각화 완료 ({len(sel)}장) → {DEBUG_VIS_DIR}\")\n",
    "print(\"✅ 전체 프로세스 완료\")\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ltb",
   "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
}