まひろ量子のハックログ

プログラミングや機械学習などの知識を記録・共有します

illust2vecを使って類似イラストを検索する

f:id:twx:20190414173122p:plain
類似画像検索

1. illust2vecとは

東北大学の齋藤真樹氏と東京大学の松井勇佑氏によって提案された「イラストの意味をベクトル化する」手法。 単語の意味をベクトル化するword2vecというものがありますが、それのイラスト版です。

ベクトル化することで、以下のようなことができるようになります。

  • 2つのベクトル間での、足し算や引き算
  • 2つのベクトルの類似度の計算

Githubでソースコードや学習済みモデルが公開されています。神かよ…

ソース github.com

モデル github.com

2. とりあえず動かしてみる

Google Colaboratory上で動かしてみます。

https://colab.research.google.com/

この記事に登場するコードは全てGoogle Colabで公開しています。 https://colab.research.google.com/drive/1uF5SgtKpluoBJLYBP0_Ta-FZMh9XChpE

まずはillust2vecをインストールします。

# illust2vecをインストール

!git clone https://github.com/rezoo/illustration2vec.git
%cd illustration2vec/
!pip install -r requirements.txt

!wget https://github.com/rezoo/illustration2vec/releases/download/v2.0.0/illust2vec_tag_ver200.caffemodel
!wget https://github.com/rezoo/illustration2vec/releases/download/v2.0.0/illust2vec_ver200.caffemodel
!wget https://github.com/rezoo/illustration2vec/releases/download/v2.0.0/tag_list.json.gz
!gzip -d tag_list.json.gz

初音ミクのイラストを入力して、タグを予測してみます。

# 初音ミクのイラストを入力とし、タグを予測してみる

import i2v
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np

illust2vec = i2v.make_i2v_with_chainer(
    "illust2vec_tag_ver200.caffemodel", "tag_list.json")
img = Image.open("images/miku.jpg")
plt.imshow(np.array(img))
illust2vec.estimate_plausible_tags([img], threshold=0.5)

f:id:twx:20190414165312p:plain
実行結果

'1girl', 'twintails', 'solo' といったタグが得られました。

3. 類似イラスト検索の仕組みを作ってみる

次に、類似イラスト検索の仕組みを作ってみようと思います。

ステップとしては、以下のようになります。

  1. 大量の画像をベクトル化したデータベースを作る
  2. 検索クエリとなる画像と、大量画像との類似度を計算する
  3. 類似度の高い画像を表示する

3.1 大量画像の準備

Google Colaboratoryから一度離れて、ローカル環境に戻ります。 ローカルPCで、Google画像検索で大量のアニメ画像を集めましょう。 google-images-download というpipモジュールを使うことで効率的に画像を集めることができます。

# pipコマンドでgoogle_images_downloadをインストール
pip install google_images_download

# google_images_downloadで画像を保存。ファイル形式はjpgのみ。
googleimagesdownload --keywords "アニメキャラ" -t jpg
googleimagesdownload --keywords "緑髪キャラ" -t jpg
googleimagesdownload --keywords "赤髪キャラ" -t jpg
googleimagesdownload --keywords "青髪キャラ" -t jpg

それぞれのフォルダに100枚、合計で400枚の画像が保存されます。たまに変な画像もあるので、そういう画像は手作業で削除してください。これらを1つのフォルダにまとめ、連番でリネームしておきます。リネームのコマンドは以下です。

ls *.jpg |  awk '{ printf "mv %s image_%03d.jpg\n", $0, NR }' | sh

f:id:twx:20190414171440p:plain
大量のイラスト画像

大量の画像が格納されたこのフォルダをillusts.zipという名前のZIPファイルにし、Google Driveにアップロードします。

Google Colaboratoryに戻り、Google Driveをマウントして先程のillusts.zipをGoogle Colaboratory上にコピーします。

from google.colab import drive
drive.mount('/content/drive')

!cp /content/drive/My\ Drive/illusts.zip ./
!unzip illusts.zip

f:id:twx:20190414171836p:plain
ColabにGoogle Driveをマウント

Colaboratory上に展開した全ての画像をベクトル化します。ここで、i2vは、イラストのパスをkeyに持ち、 イラストのベクトルをvalueに持つ辞書オブジェクトです。

from glob import glob
illusts = glob('illusts/img_*.jpg')
i2v = {}
for illust in illusts:
    img = Image.open(illust)
    v = illust2vec.extract_feature([img])[0]
    v = v / np.linalg.norm(v)
    i2v[illust] = v

3.2 検索クエリとなる画像と、大量画像との類似度を計算する

ベクトルの類似度の計算方法は色々ありますが、ここではコサイン類似度を用います。コサイン類似度とは、簡単に言うとベクトルの内積を正規化したものです。

def cos_sim(v1, v2):
    return np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2))

先程作ったi2vに格納されている全てのベクトルと、初音ミクの画像の類似度を計算していきます。ここで、i2csは、イラストのパスをkeyに持ち、 初音ミクとの類似度をvalueに持つ辞書オブジェクトです。

i2cs = {}

img = Image.open("images/miku.jpg")
target_vec = illust2vec.extract_feature([img])[0]
target_vec = target_vec / np.linalg.norm(target_vec)

for illust, vec in i2v.items():
      cs =  cos_sim(target_vec, vec)
      i2cs[illust] = cs

3.3 類似度の高い画像を表示する

コサイン類似度が高い順にソートし、上位5件を表示してみます。

sorted_i2cs = sorted(i2cs.items(), key=lambda x: -x[1])

plt.figure()
plt.imshow(np.array(img))

for illust, sim in sorted_i2cs[:5]:
    img = Image.open(illust)
    plt.figure()
    print(sim)
    plt.imshow(np.array(img))

f:id:twx:20190414173122p:plain
類似イラスト検索結果

上の1枚が検索クエリとなる画像、下の5枚が類似画像と判定された画像です。左から上位1, 2, 3, 4, 5位の順です。 確かに、似ている画像が検索できていることがわかります。

他の例も見てみましょう。

f:id:twx:20190414175044p:plain
類似イラスト検索結果その2

ピンポイントにラムちゃんを検索できたり、水着や緑髪などの特徴のあるイラストが検索できています。

以上です。この記事に登場するコードは全てGoogle Colabで公開しています。 https://colab.research.google.com/drive/1uF5SgtKpluoBJLYBP0_Ta-FZMh9XChpE

本日は「illust2vecを使って類似イラストを検索する」ことをやってみました。良い記事だと思っていただいた方は、SNSでのシェア、ブログからのリンク、「読者になる」ボタンのクリック、「★」ボタンのクリック、よろしくお願いします! ではまた!

Kozuko Mahiro's Hacklog ―― Copyright © 2018 Mahiro Kazuko