[python] train/test/validation 데이터셋 폴더에 각각 나누기

 

 

이미지 학습을 할 경우 데이터셋이 한 폴더안에 담겨 있는 경우가 많다.

 

아래와 같이 train/test/valid 데이터셋을 폴더로 각각 나누고 싶을 경우

train_folder/
  category1/
     img1.jpg
     img2.jpg
  category2/
...
...
test_folder/
  category1/
     img4.jpg
     img5.jpg
....
...

 

방법 (1) - split-folder 라이브러리 사용

 

설치

pip install split-folders
conda install split-folders

나누기

import splitfolders

# train/val/test 나누기
# 원하는 (train, validation, test 비율을 입력
splitfolders.ratio("from_folder", output="to_folder", seed=42, ratio=(.8, .1, .1))
    
# train/val 나누기
splitfolders.ratio("from_folder", output="to_folder", seed=42, ratio=(.8, .2))

(참고) https://pypi.org/project/split-folders/

split-folders

 

방법(2) - sklearn.train_test_split & shutil 사용

split-folders 라이브러리 설치가 안되어 sickit-learn에서 제공하는 train_test_split 함수로 파일을 비율에 맞게 나눠준 후에

shutil을 통해 file을 copy해주었다!

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

def move_files(file_list, source_path, destination_path):
    for file in file_list:
        image = file.split('/')[-1] + '.jpg'
        shutil.copyfile(os.path.join(source_path, image), os.path.join(destination_path, image))
        
    return
    
ls = ['Bedroom', 'Kitchen', 'Dinning', 'Livingroom']

for i in ls:
    image_files = glob(f'dataset/{i}/*.jpg')
    images = [name.replace('.jpg', '') for name in image_files]
    train_names, test_names = train_test_split(images, test_size=0.2, random_state=42, shuffle=True)

    source_dir = f'dataset/{i}/'
    train_dir = f'train_dataset/{i}/'
    test_dir = f'test_dataset/{i}/'

    move_files(train_names, source_dir, train_dir)
    move_files(test_names, source_dir, test_dir)