한줄로 Train/Test/Validation 나누기 (splitfolders : annotation도 함께 나누는법)

AI 모델 학습을 할 때 데이터셋을 train, validation, test 데이터 셋으로 나눈다.

 

간단히 설명하면,

• Train set : 학습에 사용되는 훈련용 데이터
• Test set : 학습 후에 모델의 성능을 평가하기 위해서만 사용되는 테스트용 데이터
  • 일반화 능력 평가 : 훈련집합에 없는 새로운 샘플에 대한 오류를 최소화하는 모델로 테스트 셋에 대한 높은 성능을 가지고 있는 모델
• Valid set : 모델의 일반화 능력을 높이기 위해 학습 중에 평가에 사용되는 검증 데이터 ( 예) 모의고사 )

데이터셋을 쉽게 train, validation, test 데이터 셋으로 나누는 방법을 알아보자~

 

🍯 Use splitfolders

서칭 중 split-folders라는 파이썬 라이브러리를 찾았다. 

pypi.org/project/split-folders/

 

사용방법은 매우 간단하다:

 

1) pip install split-folders /

    pip install split-folders tqdm (많은 양의 파일을 처리하는 경우, 파일 복사 프로세스를 progress bar로 시각적으로 보고 싶으면 tqdm 를 함께 설치)

 

 

2) 원하는 비율로 데이터셋을 잘라주면 된다.

Train : Validation : Test 를 8 : 1 : 1 비율로 나누는 경우에는 아래와 같이 ratio = (0.8, 0.1, 0.1) 를 해주면 된다.

Train : Validatoin / Train : Test 이렇게 2가지로만 나누는 경우에는 ratio를 튜플로 적어주면 된다 -> ratio = (0.8, 0.2)

 

데이터 경로(input folder)를 적어주고 output 디렉토리만 넣어주면 자동으로 train, val, test 폴더가 생긴다. 

import splitfolders
splitfolders.ratio('데이터 경로', output="output 폴더 경로", seed=77, ratio=(.8, 0.1,0.1))

 

*** 주의 해야할 점은 input 폴더에 class 폴더가 있어야지 작동한다. 

아래 사진 참고! class 가 없으면 그냥 input 폴더 아래에 폴더 하나 만들어서 그 안에 데이터 넣어주고 -> split 하고 파일 옮기면 된다.

 

 

3) 비율 말고도 정해진 숫자만큼만 나눌 수도 있다.

 

아래 코드는 train / test의 데이터 개수를 각각 100개씩으로 만들 때 예시이다.

 

oversample 은 기본값이 False 이다. => True 로 하면 데이터 중복 복사를 허용해 fixed 에 설정한 숫자만큼 채울 수 있다. 보통 Train 에서만 적용하고 Validation / Test 에서는 적용하지 않는다 ( 안됨.. )

# Split val/test with a fixed number of items e.g. 100 for each set.
# To only split into training and validation set, use a single number to `fixed`, i.e., `10`.
splitfolders.fixed("input_folder", output="output", seed=1337, fixed=(100, 100), oversample=False, group_prefix=None) # default values

 

 

4) Annotation 파일 (.txt, .json, .xml) 이 있는 경우 이미지 파일과 짝지어서 나눠줄 수도 있다.

 

group_prefix = 2 를 설정해주면 된다.

group_prefix = group 의 길이로 설정 (이미지, 텍스트 파일이 짝지어져 있으면 group_prefix = 2 / 이미지, 텍스트, json이 짝지어져 있으면 group_prefix = 3 등등..)

 

*** 주의 해야할 점은 데이터 경로 안에 classes.txt 같은 파일이 있으면 pair 가 없다고 에러남.. 지워주고 나중에 다시 넣어주기

splitfolders.ratio('데이터 경로', output="output 폴더 경로", seed=77, ratio=(0.8,0.1,0.1), group_prefix=2)

 

 

+) scikit-learn 에 train_test_split 써서 annotation 과 이미지 함께 train, val, test 로 나누기

splitfolders 가 있는데 train_test_split를 쓰는 이유는 ??? 

=> stratify 옵션이 있기 때문이다 

=> y 값의 (label) 의 분포를 고려해서 데이터 셋을 나눠준다 (예를 들어 이상치 탐지의 경우 0(정상) 값이 엄청 많고 1(비정상) 값이 현저히 적을 수 있는데, 분포를 고려하지 않고 데이터 셋을 나눠주면 test에 정상값만 포함될수도...)

 

아래 코드는 train_test_split 써서 train : validation : test = 8 : 1 : 1 비율로 .JPG 파일과 .json 파일을 짝지어서 함께 split 해주는 코드이다. stratify 옵션은 안썼지만 필요하면 image 데이터의 경우 DataFrame 으로 파일 경로 , class 컬럼 만들어서 class 컬럼의 분포를 기준으로 (stratify = df["class"]) 이런식으로 넣어주면 될듯하다.

import os
from glob import glob
import shutil
from sklearn.model_selection import train_test_split

#getting list of images
image_files = glob("데이터경로/*.JPG")

images = [name.replace(".JPG","") for name in image_files]

#splitting the dataset
#train:val:test = 8:1:1
train_names, test_names = train_test_split(images, test_size=0.2, random_state=42, shuffle=True)
val_names, test_names = train_test_split(test_names, test_size=0.5, random_state=42, shuffle=True)

def batch_move_files(file_list, source_path, destination_path):
    for file in file_list:
    	# 경로에서 마지막 파일명만 가져와서 확장자 붙여줌
        image = file.split('/')[-1] + '.JPG'
        txt = file.split('/')[-1] + '.json' # .txt / .json / .xml 등.. 바꿔주면됨
        shutil.copy(os.path.join(source_path, image), destination_path)
        shutil.copy(os.path.join(source_path, txt), destination_path)
    return

#data path
source_dir = "데이터경로"

#new data path
test_dir = "/ouput경로/test"
train_dir = "/ouput경로/train"
val_dir = "/ouput경로/val"
batch_move_files(train_names, source_dir, train_dir)
batch_move_files(test_names, source_dir, test_dir)
batch_move_files(val_names, source_dir, val_dir)

 

참고:

pypi.org/project/split-folders/

datascience.stackexchange.com/questions/15135/train-test-validation-set-splitting-in-sklearn

gjghks.tistory.com/78