[Python] Tensorflow로 나만의 데이터셋 만들기

안녕하세요.

오늘은 텐서플로우로 커스텀 데이터셋을 만드는 방법 예제를 진행하겠습니다.

이 예제는 분류 모델에서 사용가능한 데이터셋입니다.

 

먼저, 구글 colab을 이용해서 예제를 진행하겠습니다.

 

먼저 데이터셋에 사용할 이미지들을 다음과 같이 구성하겠습니다.

저는 위의 데이터셋이라는 폴더 안에 3개의 폴더에 이미지를 넣어놨습니다. 분류모델 학습을 위한 데이터셋이고, 각 폴더의 이름이 클래스명이 됩니다. 즉, 위의 데이터셋이라는 폴더 안에 있는 이미지들로 데이터셋을 구성하면 3개의 클래스를 가진 데이터셋이 됩니다.

 

import tensorflow as tf
from tensorflow import keras

img_gen = tf.keras.preprocessing.image.ImageDataGenerator(rescale=1./255, 
                                                          validation_split=0.2)

ImageDataGenerator 는 주로 데이터 증강을 위해 사용하는데, 지금은 증강없이 train 과 valid를 나누는 것과 rescale만 설정하겠습니다. 다만 rescale의 경우에는 종종 모델 자체 포함되 레이어에 rescale 레이어가 포함되는 경우가 있으니 꼭 확인 후에 적용하셔야 합니다. 중복 적용될 경우, 성능이 하락할 수 있습니다. validation_split은 train과 valid를 어떤 비율로 나눌지를 설정하는 것입니다. 저는 80% : 20%로 나눠서 적용하겠습니다.

 

train_set = img_gen.flow_from_directory(directory="/content/drive/MyDrive/스틱코드/데이터셋/",
                                        target_size = (480,480),
                                        class_mode='categorical',
                                        subset='training', 
                                        batch_size = 32)

위에서 생성한  imageDataGenerator에 flow_from_directory라는 폴더로 부터 이미지를 가져오는 메소드를 사용합니다. 

train set을 생성할때는  target size 는 모델에 맞는 이미지 사이즈크기로 설정하시면 되고, training 이라고 subset에 명시를 해줘야 데이터 증강을 설정했다면 정상적으로 적용이 됩니다.

 

class_mode는 "categorical", "binary", "sparse", "input", 혹은 None 중 하나를 입력해야하며 디폴트 값은 "categorical" 입니다. 

생성된 데이터셋이 가질 라벨이 어떤 분류방식을 가질지를 설정합니다.

batch_size는 사용할 배치크기를 입력하시면 됩니다.

valid_set = img_gen.flow_from_directory(directory="/content/drive/MyDrive/스틱코드/데이터셋/",
                                        target_size = (480,480),
                                        class_mode='categorical',
                                        subset='validation',
                                        shuffle = False,
                                        batch_size = 32)

train set과 마찬가지로 valid set을 구성해보겠습니다. 다만 ImageDataGenerator로 validation_split 을 설정 했다면 꼭 subset 설정값에 'validation' 이라고 명시를 해줘야 설정한 비율에 맞게 valid set 이 생성됩니다.

 

train set 과 valid set 을 합친 이미지 수가 전체 이미지 수와 맞는지 확인 한 번 해보겠습니다.

image_data = img_gen.flow_from_directory(directory="/content/drive/MyDrive/스틱코드/데이터셋/",target_size = (480,480))

 

구성된 train set의 내부 이미지 하나를 한 번 확인해 보겠습니다.

from matplotlib.image import imread
import matplotlib.pyplot as plt
import numpy as np

def get_class_string_from_index(index):
   for class_string, class_index in image_data.class_indices.items():
      if class_index == index:
         return class_string

x, y = next(train_set)
image = x[0, :, :, :]
plt.imshow((image).astype(np.uint8))
plt.axis('off')
plt.show()

 

이제 train set 과 valid set으로 모델에 학습을 진행하면 됩니다!