{ "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", "\n", "# ---------------------------\n", "# (1) 데이터셋 경로 리스트 설정\n", "# ---------------------------\n", "# 여기에 합치고 싶은 YOLO11 데이터셋 최상위 폴더 경로들을 추가하세요.\n", "DATASET_PATHS = [\n", " \"/ltb/media/ltb/90887173887158A4/Users/ltb/박스_데이터셋/박스_yolo11/it1_1819_4class_whhb\",\n", " \"/ltb/media/ltb/90887173887158A4/Users/ltb/박스_데이터셋/박스_yolo11/it1_1819_4class_whhb_albu_aug\",\n", " # \"/path/to/other_dataset1\",\n", " # \"/path/to/other_dataset2\",\n", "]\n", "\n", "# ---------------------------\n", "# (2) 출력 폴더 설정\n", "# ---------------------------\n", "# 합쳐진 train/val 파일을 저장할 최상위 경로\n", "OUTPUT_BASE = \"/ltb/media/ltb/90887173887158A4/Users/ltb/박스_데이터셋/박스_train_val/it1_1819_4class_whhb_og_albu\"\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", "\n", "# ---------------------------\n", "# (3) 분할 비율 및 랜덤 시드\n", "# ---------------------------\n", "TRAIN_RATIO = 0.8 # train : val = 0.8 : 0.2\n", "RANDOM_SEED = 42 # 재현성을 위한 시드\n", "\n", "# ---------------------------\n", "# (4) 디버깅 옵션 (시각화)\n", "# ---------------------------\n", "DEBUG_VIS = True # True이면 분할 후 일부 이미지를 시각화\n", "DEBUG_VIS_COUNT = 10 # 시각화할 이미지 개수\n", "DEBUG_VIS_DIR = os.path.join(OUTPUT_BASE, \"debug_vis\")\n", "\n", "# ---------------------------\n", "# (5) 클래스 이름 및 컬러 팔레트\n", "# ---------------------------\n", "CLASS_NAMES = {0:\"worker\",1:\"helmet\",2:\"head\",3:\"background\"}\n", "COLOR_PALETTE = [ (0,0,255),(0,255,0),(255,0,0),(0,255,255),(255,0,255),(255,255,0) ]\n", "TEXT_OFFSETS = {c:(0,-10) for c in CLASS_NAMES}\n", "\n", "print(\"✅ Cell1 설정 완료\")\n", "print(\" DATASET_PATHS:\", DATASET_PATHS)\n", "print(\" OUTPUT_BASE:\", OUTPUT_BASE)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# [Cell 2] 유틸리티 함수 정의\n", "\n", "def ensure_dir(d):\n", " \"\"\"폴더 없으면 생성\"\"\"\n", " if not os.path.exists(d):\n", " os.makedirs(d, exist_ok=True)\n", "\n", "\n", "def read_label_file(path):\n", " \"\"\"\n", " YOLO txt 파일을 읽어 클래스별 개수를 반환\n", " \"\"\"\n", " counts = defaultdict(int)\n", " total = 0\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)<1: continue\n", " cls=t[0]\n", " counts[cls]+=1\n", " total+=1\n", " except Exception as e:\n", " print(f\"❗ 라벨 읽기 오류: {path} / {e}\")\n", " return dict(counts), total\n", "\n", "\n", "def collect_all_samples(dataset_paths):\n", " \"\"\"\n", " 여러 데이터셋 경로에서 images/labels 쌍을 수집,\n", " base collision 시 경로명 prefix 추가\n", " \"\"\"\n", " samples=[]\n", " overall_counts=defaultdict(int)\n", " for ds in dataset_paths:\n", " ds_name=os.path.basename(ds.rstrip('/'))\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", " img_file = base + os.path.splitext(os.listdir(img_dir)[0])[1] # 확장자는 이미지 폴더에서 확인\n", " # 파일 매칭: 여러 확장자 지원\n", " found=None\n", " for ext in ['.jpg','.png','.jpeg']:\n", " candidate = os.path.join(img_dir, base+ext)\n", " if os.path.exists(candidate): found=candidate; break\n", " if not found:\n", " print(f\"❗ 이미지 없음: {lbl}\")\n", " continue\n", " label_path=os.path.join(lbl_dir,lbl)\n", " # 클래스 카운트\n", " cls_counts, total = read_label_file(label_path)\n", " for c,n in cls_counts.items(): overall_counts[c]+=n\n", " # 충돌 방지용 고유 이름\n", " unique_base = f\"{ds_name}_{base}\"\n", " samples.append({\n", " 'base': unique_base,\n", " 'orig_base': base,\n", " 'dataset': ds_name,\n", " 'image_path': found,\n", " 'label_path': label_path,\n", " 'class_counts': cls_counts,\n", " 'total_ann': total\n", " })\n", " return samples, overall_counts\n", "\n", "\n", "def select_test_samples(samples, n):\n", " \"\"\"각 클래스 최소 1개 + 랜덤 추가\"\"\"\n", " allc=set()\n", " for s in samples: allc.update(s['class_counts'].keys())\n", " if n