python split-folders로 이미지 데이터 분리하기 (train/val/test)

이미지 데이터셋은 보통 한 폴더 안에 몽땅 담겨 있습니다.

 

딥러닝 네트워크를 학습하고 평가 및 검증하기 위해선 이 데이터를 분리해야 하는데요. 그럴 때 유용한 python 라이브러리가 있습니다. 바로 split-folders입니다.

 

pip install split-folders
pip install split-folders[full] # progress bar를 보고 싶을때

 

사용법은 간단합니다.

split-folders를 적용하기전 데이터의 디렉터리 구조는 다음과 같아야 합니다.

input/
    class1/
        img1.jpg
        img2.jpg
        ...
    class2/
        imgWhatever.jpg
        ...
    ...

 

split-folders에는 데이터를 나누기 위한 방법이 두 가지 있습니다.

1. 비율로 나누기

import splitfolders

# train/val/test 나누기
# ratio 파라미터에 원하는 (train, validation, test) 비율을 입력합니다. ex) (0.8, 0.1, 0.1)
splitfolders.ratio("input_folder", output="output", seed=1337, ratio=(.8, .1, .1))
    
# train/val 나누기
# test는 제외하고 train, validation만 나누고 싶다면 두 개의 인자만 입력합니다. ex) (0.8, 0.1)
splitfolders.ratio("input_folder", output="output", seed=1337, ratio=(.8, .2))

 

2. 고정된 개수로 나누기

# validation만 지정해주고 싶다면 숫자 하나만 입력합니다. 나머지는 train set입니다. ex) 10
# validation/test를 지정하려면 두 개만 입력합니다. 나머지는 train set입니다. ex) (100, 100)
# train/validation/test를 모두 지정해주고 싶다면 숫자 3개를 입력합니다. ex) (300, 100, 100)
splitfolders.fixed("input_folder", output="output", seed=1337, fixed=(100, 100))

 

위 코드를 실행하면 처음 데이터 셋이 다음과 같이 변환됩니다.

output/
    train/
        class1/
            img1.jpg
            ...
        class2/
            imga.jpg
            ...
    val/
        class1/
            img2.jpg
            ...
        class2/
            imgb.jpg
            ...
    test/
        class1/
            img3.jpg
            ...
        class2/
            imgc.jpg
            ...

 

split-folders를 이용한 코드 몇 줄로 데이터를 쉽게 split 할 수 있습니다.