import os import random import shutil from collections import defaultdict # ============================================================================= # [유틸리티 함수] # ============================================================================= def ensure_dir(directory): """ 지정한 경로의 폴더가 없으면 생성합니다. 이미 존재하면 그대로 유지합니다. """ if not os.path.exists(directory): os.makedirs(directory) print(f"디렉토리 생성: {directory}") else: print(f"디렉토리 이미 존재: {directory}") def read_label_file(label_file_path): """ 주어진 라벨 파일(.txt)을 읽어 각 라인별 annotation 정보를 파싱합니다. 라벨 파일 형식: 각 라인이 "class x_center y_center width height" (공백 구분) 형태로 되어 있음. Parameters: label_file_path: 라벨 파일의 전체 경로 Returns: class_counts: {클래스 번호: 해당 파일 내 등장 횟수} 형식의 딕셔너리 total_lines: 총 annotation 라인 수 (정수) """ class_counts = defaultdict(int) total_lines = 0 try: with open(label_file_path, 'r', encoding='utf-8') as f: for line in f: line = line.strip() if not line: continue total_lines += 1 # 각 라인은 공백으로 구분되며 첫번째 값가 클래스 번호임 tokens = line.split() cls = tokens[0] class_counts[cls] += 1 except Exception as e: print(f"라벨 파일 읽기 오류: {label_file_path} / {e}") return dict(class_counts), total_lines # ============================================================================= # [데이터셋 분할 함수] # ============================================================================= def split_dataset(dataset_path, train_ratio=0.8, random_seed=42): """ 지정한 데이터셋 경로 내의 images와 labels 폴더에서 파일들을 읽어 train과 val로 분할하고, 각 그룹으로 파일들을 복사합니다. 주요 기능 및 단계: 1. dataset_path 아래의 images, labels 폴더를 확인합니다. (train, val 하위 폴더가 이미 존재하는 경우 이를 제외한 파일들만 선택합니다.) 2. images와 labels의 파일명을 매칭하여 샘플 리스트를 생성합니다. 샘플 하나는 한 이미지/라벨 파일 쌍에 해당합니다. 3. 각 라벨 파일을 읽어, 파일 내 annotation의 클래스 분포 정보를 파악합니다. 4. 전체 샘플 리스트를 랜덤하게 섞은 후, train_ratio 비율에 따라 초기 분할합니다. 5. 분할 후 val 셋에 특정 클래스가 아예 포함되지 않은 경우, train 셋에서 해당 클래스가 등장하는 파일을 하나 찾아서 val 셋으로 이동시킵니다. 6. 최종 분할된 train/val 파일들을 dataset_path 내에 아래 output 경로에 복사합니다. - images/train, images/val - labels/train, labels/val 7. 최종적으로 train/val 파일 수와 각 클래스별 annotation 갯수를 출력합니다. Parameters: dataset_path: 데이터셋 최상위 폴더 경로 (예: "/ltb/media/ltb/90887173887158A4/Users/ltb/박스_데이터셋/박스_yolo11/it1_1819_4class_whhb") train_ratio: train 데이터셋 비율 (default: 0.8) random_seed: 랜덤 시드 (default: 42) 사용 예: split_dataset("/path/to/dataset", train_ratio=0.8) """ # ------------------------------------------------------------------------- # 1. 기본 경로 및 output 경로 설정 # ------------------------------------------------------------------------- images_dir = os.path.join(dataset_path, "images") labels_dir = os.path.join(dataset_path, "labels") # output 경로: 각각 images와 labels 폴더 내에 train, val 서브폴더 생성 out_train_images_dir = os.path.join(images_dir, "train") out_val_images_dir = os.path.join(images_dir, "val") out_train_labels_dir = os.path.join(labels_dir, "train") out_val_labels_dir = os.path.join(labels_dir, "val") # 폴더 생성 (없으면 생성, 있으면 그대로 사용) ensure_dir(out_train_images_dir) ensure_dir(out_val_images_dir) ensure_dir(out_train_labels_dir) ensure_dir(out_val_labels_dir) # ------------------------------------------------------------------------- # 2. 원본 데이터(라벨, 이미지) 읽기 # ------------------------------------------------------------------------- # labels 폴더에서 train/val 폴더 제외한 .txt 파일 목록을 읽습니다. all_label_files = [] for f in os.listdir(labels_dir): full_path = os.path.join(labels_dir, f) if os.path.isfile(full_path) and f.lower().endswith(".txt"): all_label_files.append(f) if not all_label_files: print("라벨 파일을 찾을 수 없습니다.") return # images 폴더에서 train/val 폴더 제외한 파일들을 읽어, # 파일명의 베이스(확장자 제거)와 전체 파일명을 매핑하는 딕셔너리 생성 image_files_dict = {} for f in os.listdir(images_dir): full_path = os.path.join(images_dir, f) if os.path.isfile(full_path): base_name, ext = os.path.splitext(f) image_files_dict[base_name] = f # f는 이미지 파일명 (예: A_B.jpg) # ------------------------------------------------------------------------- # 3. 샘플 리스트 생성 및 라벨 파싱 # ------------------------------------------------------------------------- # 각 샘플은 { "base": 파일의 기본명, "label_path": 전체 라벨 파일 경로, # "image_path": 전체 이미지 파일 경로, "class_counts": {클래스: count} } 의 dict 형식 samples = [] overall_annotation_counts = defaultdict(int) for label_fname in all_label_files: base_name, _ = os.path.splitext(label_fname) label_path = os.path.join(labels_dir, label_fname) if base_name not in image_files_dict: print(f"이미지 파일을 찾을 수 없음 (라벨 파일과 매칭되지 않음): {label_fname}") continue image_fname = image_files_dict[base_name] image_path = os.path.join(images_dir, image_fname) # 라벨 파일 내부의 클래스 카운트 정보를 읽음 class_counts, total = read_label_file(label_path) # 전체 annotation 통계 갱신 for cls, cnt in class_counts.items(): overall_annotation_counts[cls] += cnt sample = { "base": base_name, "label_path": label_path, "image_path": image_path, "class_counts": class_counts, "total_ann": total } samples.append(sample) total_samples = len(samples) print(f"전체 샘플 수: {total_samples}") print("전체 라벨 파일 내 클래스별 annotation 갯수:") for cls, cnt in overall_annotation_counts.items(): print(f" 클래스 {cls}: {cnt}") # ------------------------------------------------------------------------- # 4. 초기 train/val 분할 (랜덤 셔플 후 train_ratio에 따라 분할) # ------------------------------------------------------------------------- random.seed(random_seed) random.shuffle(samples) train_count = int(total_samples * train_ratio) train_samples = samples[:train_count] val_samples = samples[train_count:] print(f"\n초기 분할: train = {len(train_samples)} 개, val = {len(val_samples)} 개") # ------------------------------------------------------------------------- # 5. 최소 보장 조건: 각 클래스는 val에 적어도 1개 이상은 포함 # ------------------------------------------------------------------------- # 전체에 등장하는 클래스 목록 all_classes = set() for sample in samples: all_classes.update(sample["class_counts"].keys()) # val 셋에 등장하는 클래스 목록 val_classes = set() for sample in val_samples: val_classes.update(sample["class_counts"].keys()) # 클래스별 최소 1개 이상이 val에 있도록 보장 missing_classes = all_classes - val_classes if missing_classes: print("\nval 셋에 없는 클래스가 있습니다. 해당 클래스를 포함시키기 위해 train에서 이동합니다:") for cls in missing_classes: print(f" 누락 클래스: {cls}") moved = False # train 셋에서 해당 클래스를 포함하는 샘플 찾기 for sample in train_samples: # 이미 val에 해당 클래스가 있는지 다시 확인 (혹은 다른 샘플로 보완) if cls in sample["class_counts"] and sample not in val_samples: train_samples.remove(sample) val_samples.append(sample) moved = True print(f" 파일 {sample['base']} 이동 (클래스 {cls} 포함)") # 갱신: val_classes에 추가 val_classes.update(sample["class_counts"].keys()) break if not moved: print(f" 경고: 클래스 {cls}를 포함한 파일을 train에서 찾지 못함.") else: print("\n모든 클래스가 val 셋에 최소 1개씩 포함되어 있습니다.") print(f"\n최종 분할: train = {len(train_samples)} 개, val = {len(val_samples)} 개") # ------------------------------------------------------------------------- # 6. 파일 복사: output 경로의 train, val 폴더로 이미지와 라벨 파일 복사 # ------------------------------------------------------------------------- # 복사 시, 이미 존재하는 파일은 덮어쓰지 않고 스킵함. def copy_file(src, dst): if not os.path.exists(dst): shutil.copy2(src, dst) # 복사 완료 메시지 출력 (필요시 주석 처리 가능) print(f" 복사: {src} -> {dst}") else: print(f" 스킵 (이미 존재): {dst}") print("\n[Train 셋 파일 복사 시작]") for sample in train_samples: base = sample["base"] # 파일명은 원래 이미지 파일명과 라벨 파일명 그대로 사용 (라벨은 base.txt, 이미지는 원래 파일명) image_src = sample["image_path"] label_src = sample["label_path"] # 이미지 파일명는 images 폴더 내에 그대로 복사 (하위폴더 train) image_dst = os.path.join(out_train_images_dir, os.path.basename(image_src)) label_dst = os.path.join(out_train_labels_dir, base + ".txt") copy_file(image_src, image_dst) copy_file(label_src, label_dst) print("\n[Val 셋 파일 복사 시작]") for sample in val_samples: base = sample["base"] image_src = sample["image_path"] label_src = sample["label_path"] image_dst = os.path.join(out_val_images_dir, os.path.basename(image_src)) label_dst = os.path.join(out_val_labels_dir, base + ".txt") copy_file(image_src, image_dst) copy_file(label_src, label_dst) # ------------------------------------------------------------------------- # 7. 최종 분할 통계 출력 (파일 수 및 클래스별 annotation 개수) # ------------------------------------------------------------------------- def aggregate_class_counts(samples_list): counts = defaultdict(int) for sample in samples_list: for cls, cnt in sample["class_counts"].items(): counts[cls] += cnt return dict(counts) train_class_counts = aggregate_class_counts(train_samples) val_class_counts = aggregate_class_counts(val_samples) print("\n[최종 분할 통계]") print(f"Train 파일 수: {len(train_samples)}") print(f"Val 파일 수: {len(val_samples)}") print("\nTrain 셋 클래스별 annotation 갯수:") for cls, cnt in train_class_counts.items(): print(f" 클래스 {cls}: {cnt}") print("\nVal 셋 클래스별 annotation 갯수:") for cls, cnt in val_class_counts.items(): print(f" 클래스 {cls}: {cnt}") print("\ntrain/val 분할 및 파일 복사 완료.") # ============================================================================= # [스크립트 실행] # ============================================================================= if __name__ == "__main__": # ------------------------------------------------------------------------- # 1. 데이터셋 경로와 파라미터 설정 # ------------------------------------------------------------------------- # 데이터셋의 최상위 경로 (images와 labels 폴더가 포함된 폴더) dataset_path = "/ltb/media/ltb/90887173887158A4/Users/ltb/박스_데이터셋/박스_yolo11/it1_1819_4class_whhb" # train 비율 기본값 (0.8이면 train 80%, val 20%) train_ratio = 0.8 # 필요에 따라 변경 가능 # 랜덤 시드 (재현 가능성을 위해) random_seed = 42 # ------------------------------------------------------------------------- # 2. 함수 호출하여 분할 수행 # ------------------------------------------------------------------------- split_dataset(dataset_path, train_ratio=train_ratio, random_seed=random_seed)