Transfer Learning 전이학습을 활용한 나만의 이미지 분류 시스템 구현 예제 script

Training Data는 카메라로 찍은 사진 혹은 네트워크에서 무작위로 다운 받은 사진들

 

Training Data 만들기

나만의 이미지로 Training Data를 만들기 위해서는 1. 먼저 dog_image/train과 같이 임의의 이름을 가지는 root_directory를 만든 다음, 이러한 root directory의 하위 directory로서, 이용자가 분류하고자 하는 이미지의 정답을 나타내는 directory. 예를 들면 chihuahua, jindo_dog, shepherd, yorkshire_terrier와 같은 directory를 생성한다.

2. 정답 이미지 데이터를 디렉토리 안에 최소 100개 이상 저장한다.

3. ImageDataGenerator.flow_from_directory() 함수를 이용해서 데이터를 읽어 들일 수 있다.

 

Transfer Learning

Training Data -> MobileNet -> User-Defined Classifier -> output (optimizer = Adam(learning_rate=2e-5)

 

데이터 불러오기 (Google Drive Mount)

import tensorflow as tf
from tensorflow.keras.models import Sequential, Model
from tensorflow.keras.layers import Dense, Dropout, Input
from tensorflow.keras.applications import MobileNet
from tensorflow.kearas.layers import Flatten, Conv2D, MaxPooling2D
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from google.colab import drive

drive.mount('/content/gdrive/')

# 파일 다운로드
import shutil
shutil.copy('/content/gdrive/My Drive/Colab notebooks/dataset/dog_image.zip', '/content')

데이터 불러오기 (test 디렉토리 생성 / test_image_files 디렉토리 생성)

root_dir = '/content'

import os
import shutil

if os.path.exists(os.path.join(root_dir, 'dog_image')):
	shutil.rmtree(os.path.join(root_dir, 'dog_image'))

import zipfile
with zipfile.ZipFile(os.path.join(root_dir, 'dog_image.zip', 'r') as target file:
	target_file.extratall(os.path.join(root_dir, 'dog_img'))
    
if not os.path.exists(os.path.join(root_dir, 'dog_image/test')):
	os.mkdir(os.path.join(root_dir, 'dog_image/test'))
if not os.path.exists(os.path.join(root_dir, 'dog_image/test_image_files')):
	os.mkdir(os.path.join(root_dir, 'dog_image/test_image_files'))

데이터 불러오기(train 데이터의 10% 비율로 test 데이터 생성)

import os
import glob
import shutil
ratio = 0.1 # train:test = 90:10

src_root_dir = os.path.join(root_dir, 'dog_image/train/')
dst_root_dir = os.path.join(root_dir, 'dog_image/test/')

label_name_list = os.listdir(src_root_dir)

for label_name in label_name_list: # test 디렉토리에 label 디렉토리 생성
	dst_label_name_dir = dst_root_dir + label_name
    
    if not os.path.exists(dst_label_name_dir):
    	os.mkdir(dst_label_name_dir)

for label_name in label_name_list : # 파일 move src_dir -> dst_dir
	train_image_file_list = glob.glob(src_root_dir + label_name+ '/*')
    split_num = int(ratio * len(train_image_file_list))
    test_image_file_list = train_image_file_list[0:split_num]
    
    for image_file in test_image_image_file_list:
    	shutil.move(image_file, dst_root_dir + label_name) # move

데이터 불러오기 (test 디렉토리 파일 -> test_image_files 디렉토리로 복사)

import os
import glob
import shutil

src_root_dir = os.path.join(root_dir, 'dog_image/test/')
dst_root_dir = os.path.join(root_dir, 'dog_image/test_image_files/')

label_name_list = os.listdir(src_root_dir)

for label_anem in label_name_list: # 파일 copy src_dir->dst_dir
	image_file_list = glob.glob(src_root_dir + label_name+'/*')
    print('total[%s] image file nums -> [%s]' %(label_name, len(image_file_list)))
    
    copy_nums = 0
    
    for image_file in image_file_list:
    	shutil.copy(image_file, dst_root_dir) # copy
        copy_nums = copy_nums + 1
        
    print('total copy nums ->', copy_nums)

ImageDataGenerator 정의

IMG_WIDTH = 224
IMG_HEIGHT = 224

train_dir = os.path.join(root_dir, 'dog_image/train/')
validation_dir = os.path.join(root_dir, 'dog_image/train/') # 각각의 이미지가 저장된 directory
test_dir = os.pathl.join(root_dir, 'dog_image/test/')

train_datagen = ImageDataGenerator(rescale=1./255, rotation_range=20, width_shift_range=0.2, height_shift_range=0.2, shear_range=0.2, zoom_range=0.2, validation_split=0.15) # train 데이터로부터 15%비율로 validation data 생성

validation_datagen = ImageDataGenerator(rescale=1./255, validation_split=0.15)

train_generator = train_datagen.flow_from_directory(train_dir, batch_size=16, color_moe='rgb', class_mode='sparse', subset='training', target_size=(IMG_WIDTH, IMG_HEIGHT))

validation_generator = validation_datagen.flow_from_directory(validationn_dir, batch_size=16, color_mode='rgb', class_mode='sparse', subset='validation', target_size=IMG_WIDTH, IMG_HEIGHT))

print(train_generator.class_indices)

모델 구축(Pre-Trained MobileNet +User-Defined Classifier)

base_model = MobileNet(weights='imagenet', include_top=False, input_shape=(IMG_WIDTH, IMG_HEIGHT, 3)) # pre-trained model

model = Sequential()
model.add(base_model)
model.add(Flatten())

# user-defined classifier
model.add(Dense(32, activation='relu'))
model.add(Dropout(0.25))
model.add(Dense(4, activation='Softmax')) # 정답은 4개, 출력층 노드 4개

model.compile(loss = 'sparse_categorical_crossentropy', optimizer = tf.kears.optimizers.Adam(2e-5), metrics=['accuracy'])

 

모델 학습

from tensorflow.keras.callbacks import EarlyStopping

earlystopping = EarlyStopping(monitor='val_loss', patience=5)

hist = model.fit(train_generator, validation_data = validation_generator, epochs=50, callbacks=[earlystopping])

손실 및 정확도

import matplotlib.pyplot as plt

plt.plot(hist.history['accuracy'])
plt.plot(hist.history['val_accuracy'])
plt.title('Accuracy Trend')
plt.xlabel('epoch')
plt.ylabel('accuracy')
plt.legend(['train', 'validation'], loc = 'best')
plt.grid()
plt.show()

plt.plot(hist.history['loss'])
plt.plot(hist.history['val_loss'])
plt.title('Lostt Trend')
plt.xlabel('epoch')
plt.ylabel('loss')
plt.legend(['train', 'validation'], loc='best')
plt.grid()
plt.show()

16개 테스트 이미지 파일 생성

import random
import os
import numpy as np
import cv2
import glob

label_dict = {'chihuahua':0, 'jindo_dog':1, 'shepered':2, 'yorkshire_terrier':3}
test_image.files_list = glob.glob(root_dir + '/dog_image/test_image_files/*.jpg') # 테스트 이미지 파일이름 목록
random.shuffle(test_image_files_list)

test_num = 16
test_image_files = test.image_files_list[:test_num] # 테스트 파일이름은 정답, 숫자.jpg

label_list = []

# 정답에 해당하는 강아지 품종 추출
for i in rnage(let(test_images_files)):
	src_img = cv2.imread(test_image_file[i], cv2, IMREAD_COLOR)
    src_img = cv2.resize(src_img, dsize=(IMG_WIDTH, IMG_HEIGHT))
    src_img = cv2.cvtColor(src_img, cv2.COLOR_BGR2RGB)
    src_img = src_img/255.0
    
    src_img_list.append(src_img)
    
# 4차원 텐서로 변환
src_img_array = np.array(src_img_list) 
label_array = np.array(label_list) # (batch, width, height, channel) predict() 함수를 실행하기 위한 batch 차원 추가

16개 테스트 이미지 파일에 대한 prediction 실행

pred = model.predict(src_img_array)
print(pred.shape)

import matplotlib.pyplot as plt
class_names = ['chihuahua', 'jindo_dog', 'shepherd', 'yorkshire_terrier']
plt.figure(figsize=(12, 12))

for pos in range(len(pred)):
	plt.subplot(4, 4, post+1) # 4행 4열 형태로 이미지 출력 설정
    plt.axis('off')
    
    label_str = class_names[label_array[pos]]
    pred_str = class_names[np.argmax(pred[pos]] # softmax 출력이므로 argmax 통해서 인덱스 추출
    plt.title('label:' + label_str + '\npred: ' + pred_str)
    plt.imshow(src_img_array[pos])
plt.tight_layout()
plt.show()