概要

この記事は、情報検索・検索エンジン Advent Calendar 2019の7日目の記事です。
PytorchとElasticsearchで簡単な画像(画風)検索エンジンを作りたいと思います。

という流れで、解説していきたいと思います。（今回の記事では、自分の解釈を入れながら厳密な説明を避け大まかに説明しています。論文の理解や実装について誤りがある場合は、教えて頂けると幸いです。）
コードはこちらで公開しています。

そもそものきっかけ

（少しポエムっぽいですので、手法が気になる方はこちらはスキップしてください。）
最近、なぜ脳はアートがわかるのか ―現代美術史から学ぶ脳科学入門という書籍を読みまして抽象芸術に興味を持ち始めました。
この書籍は、「科学者とアーティストが用いている還元主義的アプローチは、目的こそ互いに異なっていても（科学者は複雑な問題を解くために、また、アーティストは鑑賞者に新たな知覚的、情動的反応を喚起するために還元主義を用いる）、その方法は似通っている」というテーマのもと、抽象芸術がなぜ生まれたのか、脳はいかにして抽象芸術を理解するか、芸術への還元主義の影響などについてノーベル賞を受賞したエリック・R・カンデルさんが書いたものになります。
この本を読みながら、絵画を還元的に見るとはどういうことで、それを定式化するとしたらどういものになるだろうと考えていたところ、ディープラーニングで画風を変換する手法をふと思い出し、確かその論文内で画風を定式化していたように記憶していたので、再度その論文を読み直してみようと思ったのがきっかけになります。

また、前々からPytorchとElasticsearchのベクトル検索に興味があったので、せっかくの機会なので調べながら画風検索エンジンを作ることにしました。

画風とは

画風変換のアルゴリズムは、2016年にImage Style Transfer Using Convolutional Neural Networksというタイトルで発表されました。画像をピカソ風にしたり、ゴッホ風にしたりできます。

個人的にこちらの論文の最もすごかったことは、画風を下記の数式として表現したことにあると思っています。¹

こちらは、ディープラーニングのとある中間層における特徴量マップの内積をとったグラム行列になります。これの意味するところは、各特徴量マップが線の太さや色彩の情報をもってあり、それの相関を計算することで、この画像は線が細くて赤と緑の組み合わせが多いなどの画風に関する情報を上手く表現しています。

こちらの画風ベクトルとコンテツベクトルを利用することで画像変換を達成しています。コンテンツベクトルは、中間層の特徴量マップをそのまま使用し、画像の形状に関する情報を保持したものになります。この２つのベクトルを利用することで、元画像の形状を保ったまま画風を変換することができます。

１枚の画像から、形状に関するコンテンツベクトルと画風に関する画風ベクトルを取得することができます。この２つのベクトルをデータベースに格納することで、形状が似た画像の検索と画風が似た画像の検索ができるようになると思い、PytorchとElasticsearchで実装してみます。

Pytorchを使って、画像から画風ベクトルを抽出

公式のPytorchによる画風変換のチュートリアルを参考にして、画風ベクトルの抽出を実装します。

def gram_matrix(input):
    a, b, c, d = input.size()  # a=batch size(=1)
    # b=number of feature maps
    # (c,d)=dimensions of a f. map (N=c*d)

    features = input.view(a * b, c * d)  # resise F_XL into \hat F_XL

    G = torch.mm(features, features.t())  # compute the gram product

    # we 'normalize' the values of the gram matrix
    # by dividing by the number of element in each feature maps.
    return G.div(a * b * c * d)

def get_upper_triangle_values(value):
    n = value.size(-1)
    return value[torch.triu(torch.ones(n, n)) == 1]

def get_style_vector(image):
    return get_upper_triangle_values(gram_matrix(trained_model(image)))

まず、学習済みのモデルを使用して、中間層の特徴量マップを取得します。そしてそれのグラム行列を計算します。グラム行列は対称行列なので上三角部分の要素だけを残して、そのベクトルを保存します。

コンテンツベクトルは、中間層の特徴量をFlattenしたものになります。

def get_content_vector(image):
    trained_model(image).view(-1)

Elasticsearchにデータを格納して、似ている画風画像を検索

Elasticsearchは分散型RESTful検索/分析エンジンです。ベクトル検索を使用することで、word2vecやbertで得られたベクトルを用いて、類似文章の検索ができます。今回は画風ベクトルとコンテツベクトルを格納²して、類似画風検索を行います。
Indexの情報とPythonでのIndex作成とデータの格納は次のとおりです。

{
  "settings": {
    "number_of_shards": 2,
    "number_of_replicas": 1
  },
  "mappings": {
    "dynamic": "true",
    "_source": {
      "enabled": "true"
    },
    "properties": {
      "title": {
        "type": "text"
      },
      "text": {
        "type": "text"
      },
      "content_vector": {
        "type": "dense_vector",
        "dims": 100
      },
      "style_vector": {
        "type": "dense_vector",
        "dims": 100
      }
    }
  }
}

from elasticsearch.helpers import bulk

# create index
index_name = "image"
index_file = "index.json"
client = Elasticsearch()
client.indices.delete(index=index_name, ignore=[404])
with open(index_file) as index_file:
    source = index_file.read().strip()
    client.indices.create(index=index_name, body=source)


def create_document(title, content_vector, style_vector):
    return {
        '_op_type': 'index',
        '_index': index_name,
        'title': title,
        'content_vector': content_vector,
        'style_vector': style_vector
    }

# insert data
bulk(client, documents)

類似画風検索をするには、入力画像の画風ベクトルを計算して、そのベクトルをElasticsearchに渡すことで、類似画像を得ることができます。

def search(vector_name, vector, topn=3):
    script_query = {
        "script_score": {
            "query": {"match_all": {}},
            "script": {
                "source": "cosineSimilarity(params.query_vector, doc['{}']) + 1.0".format(vector_name),
                "params": {"query_vector": vector}
            }
        }
    }
    response = client.search(
        index=index_name,
        body={
            "size": topn,
            "query": script_query,
            "_source": {"includes": ["title"]}
        }
    )
    return response

i = 9
query_vector = reduced_style_vectors[i, :].tolist()
pprint(search("style_vector", query_vector, 5))
query_vector = reduced_content_vectors[i, :].tolist()
pprint(search("content_vector", query_vector, 5))

結果

入力画像（ルノワール）

検索結果	類似コンテンツ画像	類似画風画像
1
2
3
4

類似画風画像の１、２位は、入力画像を描いたルノワールによって書かれた作品になります。類似画風検索のほうがより印象派の画風に似た結果になっているように思います。
(この記事では、手元でいくつかの画像を検索することで定性的に評価をしました。定量的にする場合は、画像に印象派やキュビズムなどのタグを付けて、タグの一致率などで検証できますが、今回は時間とデータの理由で定量的には評価はできておりません。。。)

(おまけ)Kibanaでデータ確認

下記のdocker-compose.yamlを用意して、docker-compose upでKibanaとElasticsearchを起動します。

version: '3.7'
services:

  elasticsearch:
    image: docker.elastic.co/elasticsearch/elasticsearch:7.3.2
    ports:
      - 9200:9200
    volumes:
      - es-data:/usr/share/elasticsearch/data
    tty: true
    environment:
      discovery.type: single-node

  kibana:
    image: docker.elastic.co/kibana/kibana:7.3.2
    ports:
      - 5601:5601
    depends_on:
      - elasticsearch

volumes:
  es-data:
    driver: local

http://localhost:5601/にアクセスすると格納したデータを確認することができます。

今後(今回できなかったこと。。)

metric learningを利用してより画風や形状を表現したベクトルの利用
style vectorを計算するさいにどの中間層を使うかで結果が変わってくるので、違う中間層を試して傾向を見てみたい（印象派などの情報は後半の層でよりきれいに表現されているなど）
次元圧縮手法と距離計算の評価
Elasticsearchで近似近傍探索の実装

　参考資料

画像の質感を表現するものとして、2015年に提案されています。Texture Synthesis Using Convolutional Neural Networks ↩
今回得られる画風ベクトルとコンテツベクトルは、それぞれ8256次元と8192次元になりますが、Elastic searchのベクトル検索の上限値は1000のため格納できません。そこで次元圧縮して100次元にしてからデータを格納しました。 ↩

PytorchとElasticsearchで画風検索

概要

目次