Kerasで画像分類

Kerasを使用した画像判定プログラムを作成してみました。元ネタはkill_me_learningです。

.

使用方法

Kerasが動作可能な環境を用意して、以下の様なフォルダ構造を用意します。

train_dist
train_data-*
    illust_1
    illust_2
    ...
nanoha_predict.py
nanoha_train.py

train_data-illustフォルダ内に、分類したい画像をフォルダ毎にわけて保存しておきます。

例えば、リリカルなのはシリーズの「なのは」「はやて」「フェイト」を分類しようと思ったら、

train_data-nanoha
    fate
    hayate
    nanoha

といった感じのフォルダを構成して、それぞれのフォルダ内に画像を格納していきます。

学習方法

$ python nanoha_train.py {学習対象のフォルダ名称}

# 上の例(なのはキャラクターの分類)でしたら、以下の様に指定します。
$ python nanoha_train.py train_data-nanoha

実行すると学習を開始し、120エポックした後に学習結果をtrain_distに保存します。

判定方法

$ python nanoha_predict.py {学習対象のフォルダ名称} {判定させたい画像}

実行すると学習結果からモデルを読み込み、対象画像がどのクラスに該当するかを表示します。

.

Kerasのプログラム

学習用

# -*- coding: utf-8 -*-
"""
    @brief クラス分類アプリケーション
"""
from __future__ import print_function
import sys
import json

import keras.preprocessing.image
from keras.optimizers import RMSprop
from keras.models import Sequential
from keras.layers import Conv2D, Convolution2D, Activation, Dropout, MaxPooling2D, Dense


IMAGE_W = 150
IMAGE_H = 150
COLOR_MODE = 3

COLOR_MODE_NAME = {
    1: "grayscale",
    3: "rgb"
}


# ============================================================================
def create_model(train_shape, class_size):
    """モデルの作成
    """

    o_model = Sequential()

    o_model.add(Conv2D(32, (3, 3), padding="valid", input_shape=train_shape))
    o_model.add(Activation("relu"))
    o_model.add(Conv2D(32, (3, 3)))
    o_model.add(Activation("relu"))
    o_model.add(MaxPooling2D(pool_size=(2, 2)))
    o_model.add(Dropout(0.25))

    o_model.add(Conv2D(64, (3, 3), padding="valid"))
    o_model.add(Activation("relu"))
    o_model.add(Conv2D(64, (3, 3)))
    o_model.add(Activation("relu"))
    o_model.add(MaxPooling2D(pool_size=(2, 2)))
    o_model.add(Dropout(0.25))

    o_model.add(keras.layers.Flatten())
    o_model.add(Dense(256))
    o_model.add(Activation("relu"))
    o_model.add(Dropout(0.5))
    o_model.add(Dense(class_size))
    o_model.add(Activation("softmax"))

    rmsplop = RMSprop(lr=0.001, rho=0.9, epsilon=1e-08, decay=0.0)
    o_model.compile(loss="categorical_crossentropy", optimizer=rmsplop, metrics=["accuracy"])

    return o_model


# ============================================================================
def main():
    """プログラムエントリー
    """

    train_datagen = keras.preprocessing.image.ImageDataGenerator(
        rescale=1.0 / 255,
        zoom_range=0.0,
        horizontal_flip=False,
        vertical_flip=False
    )

    train_generator = train_datagen.flow_from_directory(
        sys.argv[1],
        target_size=(IMAGE_W, IMAGE_H),
        color_mode=COLOR_MODE_NAME[COLOR_MODE],
        batch_size=32,
        class_mode="categorical"
    )

    with open("./" + sys.argv[1] + "/class_information.json", "wb") as h_writer:
        json.dump(
            {
                "train_shape": {"image_w": IMAGE_W, "image_h": IMAGE_H, "color_mode": COLOR_MODE},
                "class_indices" : train_generator.class_indices
            },
            h_writer
        )

    class_size = len(train_generator.class_indices)

    o_model = create_model(
        (IMAGE_W, IMAGE_H, COLOR_MODE),
        class_size
    )

    nb_epoch = 120
    o_model.fit_generator(
        train_generator,
        steps_per_epoch=3,
        epochs=nb_epoch,
        validation_steps=120
    )

    # save
    o_model.save("./train_dist/" + sys.argv[1] + ".h5")


if __name__ == "__main__":
    main()



# ---------------------------------------------------------------------- [EOF]

 

判定用

# -*- coding: utf-8 -*-
"""
    @brief クラス分類アプリケーション
"""
from __future__ import print_function
import sys
import json

import keras
import keras.preprocessing.image

import PIL


# ============================================================================
def main():
    """プログラムエントリー
    """

    with open("./" + sys.argv[1] + "/class_information.json", "rb") as h_reader:
        dict_information = json.load(h_reader)
        train_shape = dict_information["train_shape"]

        dict_class = {}
        for name, idx in dict_information["class_indices"].items():
            dict_class[idx] = name

    # load
    o_model = keras.models.load_model("./train_dist/" + sys.argv[1] + ".h5")

    image_pathname = sys.argv[2]
    img = keras.preprocessing.image.load_img(image_pathname)
    newimage = img.resize((train_shape["image_w"], train_shape["image_h"]))
    if train_shape["color_mode"] == 1:
        newimage = PIL.ImageOps.grayscale(newimage)

    newsize = keras.preprocessing.image.img_to_array(newimage)
    newsize = newsize.reshape((1,) + newsize.shape)
    newsize /= 255


    preds = o_model.predict_on_batch(newsize)

    print(image_pathname)
    list_result = []
    for idx in range(len(dict_class)):
        list_result.append(
            u"%016s\t%7.03f" % (dict_class[idx], preds[0][idx])
        )

    list_result.sort(key=lambda v: float(v.split()[1].strip()))
    list_result.reverse()

    for record in list_result:
        print(record)


if __name__ == "__main__":
    main()



# ---------------------------------------------------------------------- [EOF]

.

学習させてみたもの

Twitterにも掲載しましたが、イラストを描いている方の絵を学習させて、同じモチーフの絵を与えた時に誰の絵として判定されるかを試して見ました。

本人の絵(一番上)が本人の絵だと認識されないのが、ちょっと悲しいw