[TensorFlow] CNNs을 이용한 이미지 분류

CNNs을 이용하여 이미지를 분류해 보겠습니다.

이미지 분류에 대한 내용은 이곳을 확인해 주세요.

 

이미지 분류 class는 airplane, car, cat, dog, flower, fruit, motorbike, person으로 분류 개수는 총 8개입니다.

이미지는 CIFAR나 다른 데이터셋을 사용해도 됩니다.

 

또한 학습에 사용할 이미지를 /data/train과 /data/test 디렉토리에 나누어 저장합니다. /data/train에는 이미지 class별로 디렉토리를 생성합니다.

 

모델의 구성은 다음과 같습니다.

 

 

학습에 사용할 optimizer는 Adam이며, softmax를 이용한 multi class 분류이기 때문에 loss는 sparse_categorical_crossentropy를 사용합니다.

 

학습 성능(Accuracy, Loss)은 다음과 같습니다.

이곳에서 실험한 Pre-trained 모델의 정확도는 100%였으나, 직접 설계한 모델은 검증 데이터에 대한 정확도가 82%정도입니다.

 

테스트 데이터를 모델의 입력으로 사용한 결과에서도 잘못 분류한 데이터가 종종 보입니다.

 

아래는 전체 코드입니다.

 

[train.py]

# Import library
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder
from pathlib import Path
import cv2
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import click
import pickle
import tensorflow as tf
from tensorflow.keras.optimizers import Adam


def generate_data_frame(data_dir):
    dir = Path(data_dir)
    filepaths = list(dir.glob(r'**/*.jpg'))

    labels = [str(filepaths[i]).split("\\")[-2] \
              for i in range(len(filepaths))]

    filepath = pd.Series(filepaths, name='Filepath').astype(str)
    labels = pd.Series(labels, name='Label')

    df = pd.concat([filepath, labels], axis=1)
    df = df.sample(frac=1, random_state=0).reset_index(drop=True)

    return df


def image_data_generator(image_paths, labels):
    image_list = []
    label_list = []

    for i in range(len(image_paths)):
        image = cv2.imread(image_paths[i], cv2.IMREAD_COLOR)
        image = cv2.resize(image, (227, 227))
        label = labels[i]

        image_list.append(image)
        label_list.append(label)

    return np.stack(image_list), np.array(label_list)


def my_model(class_num):
    input_shape = (227, 227, 3)
    x = tf.keras.Input(shape=input_shape)

    conv1 = tf.keras.layers.Conv2D(filters=96, kernel_size=11, activation='relu', strides=4)(x)
    pool1 = tf.keras.layers.MaxPooling2D((3, 3), strides=2)(conv1)  # overlapped pooling

    conv2 = tf.keras.layers.Conv2D(filters=256, kernel_size=5, activation='relu', strides=1, padding='same')(pool1)
    pool2 = tf.keras.layers.MaxPooling2D((3, 3), strides=2)(conv2)

    conv3 = tf.keras.layers.Conv2D(filters=384, kernel_size=3, activation='relu', strides=1, padding='same')(pool2)
    conv4 = tf.keras.layers.Conv2D(filters=384, kernel_size=3, activation='relu', strides=1, padding='same')(conv3)
    conv5 = tf.keras.layers.Conv2D(filters=256, kernel_size=3, activation='relu', strides=1, padding='same')(conv4)
    pool3 = tf.keras.layers.MaxPooling2D((3, 3), strides=2)(conv5)

    # FC
    f = tf.keras.layers.Flatten()(pool3)
    f = tf.keras.layers.Dropout(0.5)(f)
    f = tf.keras.layers.Dense(4096, activation='relu')(f)
    f = tf.keras.layers.Dropout(0.5)(f)
    f = tf.keras.layers.Dense(4096, activation='relu')(f)
    out = tf.keras.layers.Dense(class_num, activation='softmax')(f)

    model = tf.keras.Model(inputs=x, outputs=out)

    return model


@click.command()
@click.option('--data_dir', default='data/train', help='Data path')
@click.option('--batch_size', default=128, help='Batch size')
@click.option('--epochs', default=20, help='Epochs')
@click.option('--model_name', default='tensorflow_model', help='Model name')
def run(data_dir, batch_size, epochs, model_name):
    # Load data
    df = generate_data_frame(data_dir)
    print('========== Data shape ==========')
    print(df.shape)
    print('========== Data ================')
    print(df)
    print()

    # Label encoding
    class_label = LabelEncoder()
    df['Label'] = class_label.fit_transform(df['Label'].values)
    print('========== Class ================')
    print(np.sort(df['Label'].unique()))
    print('========== Class number =========')
    class_num = len(df['Label'].unique())
    print(class_num)

    # Separate dataset
    train_df, valid_df = train_test_split(df, test_size=0.2, random_state=0)
    X_train, y_train = image_data_generator(train_df['Filepath'].tolist(), train_df['Label'].tolist())
    X_valid, y_valid = image_data_generator(valid_df['Filepath'].tolist(), valid_df['Label'].tolist())

    # Define model
    model = my_model(class_num=class_num)
    model.summary()
    optimizer = Adam(lr=1e-3)
    model.compile(loss='sparse_categorical_crossentropy', optimizer=optimizer, metrics=['accuracy'])

    # Train model
    history = model.fit(X_train, y_train, validation_data=(X_valid, y_valid), batch_size=batch_size, epochs=epochs)

    pd.DataFrame(history.history)[['accuracy', 'val_accuracy']].plot()
    plt.title("Accuracy")
    plt.show()

    pd.DataFrame(history.history)[['loss', 'val_loss']].plot()
    plt.title("Loss")
    plt.show()

    model.save(model_name + '.h5')

    # Map the label
    mapping = dict(zip(class_label.classes_, range(len(class_label.classes_))))
    labels = (mapping)
    labels = dict((v, k) for k, v in labels.items())
    with open('labels.pkl', 'wb') as f:
        pickle.dump(labels, f)


if __name__ == '__main__':
    run()

 

[test.py]

import click
import numpy as np
import pandas as pd
import pickle
import os
import cv2
from tensorflow.python.keras.saving.save import load_model


def load_data(images_dir):
    name_list = []
    image_list = []

    files = os.listdir(images_dir)

    for file in files:
        try:
            path = images_dir + '/' + file

            image = cv2.imread(path, cv2.IMREAD_COLOR)
            image = cv2.resize(image, (227, 227))

            name_list.append(file)
            image_list.append(image)

        except FileNotFoundError as e:
            print('ERROR : ', e)

    return np.array(name_list), np.stack(image_list)


@click.command()
@click.option('--data_dir', default='data/test', help='Data path')
@click.option('--model_name', default='tensorflow_model', help='Model name')
def run(data_dir, model_name):
    image_names, X_test = load_data(data_dir)

    loaded_model = load_model(model_name + '.h5')

    pred = loaded_model.predict(X_test)
    pred = np.argmax(pred, axis=1)

    with open('labels.pkl', 'rb') as f:
        labels = pickle.load(f)
        pred = [labels[k] for k in pred]

    results_df = pd.DataFrame({'image_names': image_names, 'class': pred})
    results_df.to_csv('tensorflow_results.csv', index=False)


if __name__ == '__main__':
    run()