[Tensorflow] 김치 이미지 분류기 - 데이터 분류

AI HUB에서 가져온 데이터를 확인해보겠습니다. 김치 압축 폴더의 압축을 해제하면 다음의 디렉터리 구조가 나오게 됩니다.

 

김치/
  갓김치/
  나박김치/
  무생채/
  배추김치/
  백김치/
  부추김치/
  열무김치/
  오이소박이/
  총각김치/
  파김치/

 

각각의 폴더에는 1000장의 이미지가 있지만. jpg 이미지만 사용할 것이므로 다른 포맷을 가진 데이터를 삭제하였습니다. 검증 분할을 사용하기 위해서 총 이미지 데이터에서 70%는 훈련에 나머지 30%는 검증에 사용할 수 있도록 이미지를 분리하겠습니다.

 

import os
import shutil

 

필요한 라이브러리를 import 해줍니다.

 

directory_list = ['./train', './valid']

for directory in directory_list:
    if not os.path.isdir(directory):
        os.makedirs(directory)

 

 

훈련과 검증 데이터를 담을 폴더가 필요하므로 현재 디렉토리에 train과 valid 폴더가 없다면 생성해줍니다.

 

 

def make_query_folder_in_directory(query):
  for directory in directory_list:
    if not os.path.isdir(directory + '/' + query):
      os.makedirs(directory + '/' + query)

 

마찬가지로 훈련과 검증 폴더 안에 각각의 데이터를 저장할 폴더가 필요하므로 (test/갓김치, valid/갓김치) 해당 디렉터리에 폴더가 없다면 생성시켜줍니다.

 

def train_valid_dataset_split(query, valid_scale=0.3):
    file_names = os.listdir(query)
    valid_count = int(len(file_names) * valid_scale)

    count = 0
    for file_name in file_names:
      if count > valid_count:
        print(f'[Train Dataset] {query} {file_name}')
        shutil.move(query + '/' + file_name, './train/' + query + '/' + file_name)
      else:
        print(f'[Valid Dataset] {query} {file_name}')
        shutil.move(query + '/' + file_name, './test/' + query + '/' + file_name)
      count += 1

 

여기서 valid_scale은 전체 데이터에서 검증에 사용할 데이터의 비율을 저장하는 변수로 0~1 사이의 실수 값을 입력할 수 있습니다. 검증과 훈련에 사용될 이미지의 수가 정해졌으므로 반복문을 통해서 각각의 데이터가 올바른 폴더에 갈 수 있도록 해줍니다.

 

file_list = ['파김치','백김치','부추김치','열무김치','총각김치','무생채','갓김치','배추김치','나박김치','오이소박이']

for file_name in file_list:
  make_query_folder_in_directory(file_name)
  train_valid_dataset_split(file_name)
  shutil.rmtree(file_name)

 

file_list에서 하나씩 정보를 가져와서 앞서 정의했던 함수들을 실행해줍니다. 학습과 검증 데이터까지 분리가 되었다면 shutil.rmtree(file_name) 함수를 실행시켜주어 기존에 이미지 데이터가 담겨있었던 폴더를 삭제해줍니다.