二次元イラストの特徴を検出する[illustration2vec]

#1.illustration2vecとは？
illustration2vecは事前学習済みのニューラルネットワークで、
DanbooruとSafebooruから100万枚の画像を用いて学習したもの
タグを正解ラベルとしているので、簡単に二次元イラストの特徴を検出できるので便利
VGGを参考にしたモデルで学習してるらしい詳しくは論文を参照
Caffe版とChainer版があるが今回はChainerの方を使用する。
#2.実際に検出する
https://github.com/rezoo/illustration2vec からリポジトリを入手、i2vが本体
get_models.shにあるリンクから予めtag_list.json.gzとillust2vec_tag_ver200.caffemodelを入手
実行するコードと同じディレクトリに入れる。
(illust2vec_tag.prototxtとillust2vec_ver200.caffemodelはFeatureVectorを検出する用のものなので今回は使わない)

illustration2vecにはimagesにミクさんの画像が入っているが今回は使わずにフリー素材のhttp://www.jewel-s.jp/download/
からf353.pngを使う。

実際に検出する

main.py

import i2v
from PIL import Image

def estimate_one(img):
    illust2vec = i2v.make_i2v_with_chainer("illust2vec_tag_ver200.caffemodel", "tag_list.json")
    img = Image.open(img)
    tag = illust2vec.estimate_plausible_tags([img])
    return tag[0]

def main():
    tag = estimate_one('images/f353.png')
    print(tag)

if __name__ == "__main__":
    main()

注意するのはestimate_plausible_tagsに渡す引数は画像のリストで、戻り値は辞書のリストということ
keyはcharacter(キャラ名)、copyright(著作)、general(特徴)、rating(18禁かどうか)

結果

{'copyright': [], 'character': [], 'rating': [('safe', 0.9474104642868042), ('questionable', 0.04457685351371765), ('explicit', 0.00435987114906311)], 
'general': [('1girl', 0.9905315041542053), ('blonde hair',0.9727756977081299), ('green eyes', 0.9559593796730042), ('solo', 0.8911849856376648), 
('open mouth', 0.6233431696891785), ('animal ears', 0.5574820637702942), ('smile', 0.4673902988433838), ('long hair', 0.41225337982177734), 
('fang', 0.3859747648239136), ('dress', 0.3833463191986084), ('frills', 0.3065197467803955), ('bow', 0.28883492946624756), ('ribbon', 0.2731117606163025), 
('fangs', 0.2701604962348938)]}

大体あってる。
#3.問題点
基本的なやり方は上記だけで大丈夫、ただ大量に画像を処理する場合この章を読んでもらいたい

実際に実行してみるとわかると思うが、案外計算に時間が掛かる
実はこれ計算に時間が掛かっているというよりCaffeFunctionという
CaffeのモデルをChainerのモデルにコンバート(変換)するのに掛かっている時間
これのなにが問題かというと大量に画像をタグ付けする場合、そのリストを渡すと
OutOfMemory(メモリ不足)になる可能性があるので上記のestimate_one関数をループさせる方が安全なのだが
ループする都度コンバートしてしまう。ゆえに予めループさせる前にコンバートさせておく
と大幅に高速化する。
なのでchainer_i2vにあるmake_i2v_with_chainerを少し改良する。

chainer_i2v.py

def make_i2v_with_chainer(param_path, tag_path=None, threshold_path=None, p_net=None):
    if p_net is not None:
        net = p_net
    else:
        # ignore UserWarnings from chainer
        with warnings.catch_warnings():
            warnings.simplefilter('ignore')
            net = CaffeFunction(param_path)

    kwargs = {}
    if tag_path is not None:
        tags = json.loads(open(tag_path, 'r').read())
        assert(len(tags) == 1539)
        kwargs['tags'] = tags

    if threshold_path is not None:
        fscore_threshold = np.load(threshold_path)['threshold']
        kwargs['threshold'] = fscore_threshold

    return ChainerI2V(net, **kwargs)

main.py

import i2v
from PIL import Image
from utility import utility
import warnings
from chainer.links.caffe import CaffeFunction

def estimate_one(img, param_path, tag_path, p_net):
    illust2vec = i2v.make_i2v_with_chainer(param_path=param_path, tag_path=tag_path, p_net=p_net)
    img = Image.open(img)
    tag = illust2vec.estimate_plausible_tags([img])
    return tag[0]

def main():
    tags = []
    param_path = "illust2vec_tag_ver200.caffemodel"
    tag_path = "tag_list.json"

    with warnings.catch_warnings():
        warnings.simplefilter('ignore')
        print("start converting...")
        p_net = CaffeFunction(param_path)
        print("end converting")

    for img in utility.find_all_files('images/'):
        tag = estimate_one(img, param_path, tag_path, p_net)
        tags.append(tag)

    print(tags)

if __name__ == "__main__":
    main()

find_all_filesはimagesにあるファイルをyieldですべて返す自作のジェネレーター

#4.まとめ
正直illustration2vecだけで画像をフォルダ分けするアプリとかできそうなので非常に便利
ただillustration2vecは計算がCPUのみなので大量の画像分類には少々時間が掛かるため
GPU化できれば次の記事にしたい。

著作権
Copyright (c) 2015 Masaki Saito and Yusuke Matsui.
under the MIT License