파일을 Train/Test 디렉토리로 나눠서 저장하는 방법

 

import random
import tensorflow as tf
from tensorflow.keras.optimizers import RMSprop
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from shutil import copyfile

수집한 이미지 파일들을, ImageDataGenerator를 이용해 학습데이터(numpy array)로 변환하기 위해,

그 형식에 맞게 디렉토리를 만들고 저장하는 법을 알아보자.

 

 

1. 연습용 데이터 다운로드 후 압축해제

# 연습용 이미지파일
# /tmp 경로에 압축파일을 다운로드 받고 압축 푼다.
import zipfile

!wget --no-check-certificate \
    'https://block-edu-test.s3.ap-northeast-2.amazonaws.com/kagglecatsanddogs_5340.zip' \
    -O "/tmp/cats-and-dogs.zip"

local_zip = '/tmp/cats-and-dogs.zip'
zip_ref = zipfile.ZipFile(local_zip, 'r')
zip_ref.extractall('/tmp')
zip_ref.close()

압축 풀린 상태

 

 

2. 파일을 저장할 디렉토리 생성

 

/tmp 디렉토리 안에, 학습을 위한 데이터를 분류하기 위해
cats-v-dogs 디렉토리를 만들고, 
그 아래 training 과 testing 디렉토리 만든 후
각각 디렉토리 안에  cats 와 dogs 디렉토리를 만든다.

 

import os

try:
  os.mkdir('/tmp/cats-v-dogs')

  os.mkdir('/tmp/cats-v-dogs/training')
  os.mkdir('/tmp/cats-v-dogs/testing')

  os.mkdir('/tmp/cats-v-dogs/training/cats')
  os.mkdir('/tmp/cats-v-dogs/training/dogs')

  os.mkdir('/tmp/cats-v-dogs/testing/cats')
  os.mkdir('/tmp/cats-v-dogs/testing/dogs')


except OSError:
  pass

 

 

3. Train, Test 데이터 분리

 

함수 정의

1. 소스 경로에 있는 파일들을 잘 섞고,
2. SPLIT_SIZE 만큼 Train, Test 데이터를 분리한 후,
3. training 디렉토리와 testing 디렉토리에 파일을 저장 시켜 주는 함수 정의 (파일크기가 0보다 큰 파일만 저장)
4. 파라미터
   
   • SOURCE : 원본 데이터 파일의 경로
   • TRAINING : Train 데이터 파일의 경로
   • TESTING : Test 데이터 파일의 경로
   • SPLIT_SIZE : 데이터를 분리 할 수치 (Train셋의 비율)

import random
from shutil import copyfile

def split_data(SOURCE, TRAINING, TESTING, SPLIT_SIZE):

  # 원본 경로의 파일명들을 모두 가져온다.
  file_names = os.listdir( SOURCE  )

  # 가져온 파일명을 잘 섞어준다.
  shuffled_files = random.sample( file_names  , len(file_names) )

  # 학습용과 테스트용 분리를 위해 인덱스를 구한다.
  index = int(  len(shuffled_files)  *   SPLIT_SIZE )

  # 섞인 파일명 리스트에서, 해당 인덱스만큼 잘라서, 학습용과 테스트용 파일명 분리
  training_images = shuffled_files[0 : index]
  test_images = shuffled_files[ index :  ]

  # 학습용과 테스트용으로 각각 해당 폴더로, 파일을 카피한다.
  for file_name in training_images :
    if os.path.getsize( SOURCE + file_name  ) > 0 :
      copyfile( SOURCE + file_name  , TRAINING + file_name )

  for file_name in test_images : 
    if os.path.getsize( SOURCE + file_name  ) > 0 :
      copyfile( SOURCE + file_name, TESTING + file_name )

 

 

함수 호출

 

CAT_SOURCE_DIR = "/tmp/PetImages/Cat/"
TRAINING_CATS_DIR = "/tmp/cats-v-dogs/training/cats/"
TESTING_CATS_DIR = "/tmp/cats-v-dogs/testing/cats/"
DOG_SOURCE_DIR = "/tmp/PetImages/Dog/"
TRAINING_DOGS_DIR = "/tmp/cats-v-dogs/training/dogs/"
TESTING_DOGS_DIR = "/tmp/cats-v-dogs/testing/dogs/"
split_size = 0.9

split_data(CAT_SOURCE_DIR, TRAINING_CATS_DIR, TESTING_CATS_DIR, split_size)
split_data(DOG_SOURCE_DIR, TRAINING_DOGS_DIR, TESTING_DOGS_DIR, split_size)

 

 

결과 확인

print(len(os.listdir('/tmp/cats-v-dogs/training/cats/')))
print(len(os.listdir('/tmp/cats-v-dogs/training/dogs/')))
print(len(os.listdir('/tmp/cats-v-dogs/testing/cats/')))
print(len(os.listdir('/tmp/cats-v-dogs/testing/dogs/')))