import os
import pandas as pd
from sklearn.model_selection import train_test_split

# 이미지 파일이 들어있는 디렉토리 경로
image_directory = "path/to/your/image/directory"

# 클래스별 디렉토리명과 클래스명 매핑
class_directories = ["class1", "class2", "class3"]  # 클래스 디렉토리명 수정
class_labels = ["label1", "label2", "label3"]  # 클래스명 수정

# 이미지 파일과 클래스 레이블 매핑을 저장할 리스트
data = []

# 클래스별 이미지 파일 경로 수집
for i, class_dir in enumerate(class_directories):
    image_paths = os.listdir(os.path.join(image_directory, class_dir))
    for image_path in image_paths:
        image_path_full = os.path.join(image_directory, class_dir, image_path)
        data.append((image_path_full, class_labels[i]))

# 데이터를 DataFrame으로 변환
df = pd.DataFrame(data, columns=["image_path", "label"])

# 학습 및 검증 데이터셋으로 분할
train_df, val_df = train_test_split(df, test_size=0.2, random_state=42)

# CSV 파일로 저장
train_df.to_csv("train_dataset.csv", index=False)
val_df.to_csv("val_dataset.csv", index=False)