# Dokcerイメージをpullする
docker pull docker.elastic.co/elasticsearch/elasticsearch:7.3.2

# ポート開けたり必要な環境変数を設定してDocke起動する
docker run -p 9200:9200 -p 9300:9300 -e "discovery.type=single-node" docker.elastic.co/elasticsearch/elasticsearch:7.3.2

import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from torchvision.datasets import FashionMNIST

# データセットをダウンロード
fmnist_data = FashionMNIST('./data/fashion-mnist', train=True, download=True, transform=transforms.ToTensor())
data_loader = DataLoader(fmnist_data, batch_size=4, shuffle=False)
data = data_loader.dataset

# Elasticsearchにインデクシング
es = Elasticsearch("http://localhost:9200/")

# インデックスの作成と、マッピングの設定

mapping = {
    "images": {
        "properties": {
            # 画像のインデックス
            "image_index":{
                "type": "integer"
            },
            # targetのクラスを入れておく（入れなくても良い）
            "category": {
                "type": "integer"
            },
            # Fashion-MNISTは28*28の画像なので784次元
            # ちなみに、公式でサポートしている上限は1024
            "image_vector": {
                "type": "dense_vector",
                "dims": 784
            }
        }
    }
}

# settingsを指定してインデックスを作成
es.indices.create(index='raw-images')

# 作成したインデックスのマッピングを指定
es.indices.put_mapping(index='raw-images', doc_type='images', body=mapping, include_type_name=True)

# インデクシング用の関数
def _load_data(index_, type_, data, batch_num):
    for i, row in enumerate(data):
        image, label = row
        body = {
            "image_index": batch_num + i,
            "category": int(label),
            # ここややこしいけれど普通のfloatに変換しているだけ
            # numpy系はそのまま突っ込めないっぽい
            "image_vector": list(map(lambda x: x.item(), image.numpy().flatten()))
        }
        yield {"_index": index_, "_type": type_, "_source": body}

# インデクシング（60000件は1分かからないくらいでインデクシング終了）
batch_size = 1000
dataset_with_class = list(zip(data.data, data.targets))
for i in range(math.ceil(len(dataset_with_class) / batch_size)):
    slice_data = dataset_with_class[i*batch_size:i*batch_size + batch_size]
    helpers.bulk(es, _load_data("raw-images", "images", slice_data, i*batch_size))

input_image_index = 0
image, label = dataset_with_class[input_image_index]

# 検索
query_vector = list(map(lambda x: x.item(), image.numpy().flatten()))

# 画像インデックス取得
body = {
    "query": {
        "script_score": {
            "query": {"match_all": {}},
            "script": {
                "source": "cosineSimilarity(params.query_vector, doc['image_vector']) + 1.0",
                "params": {"query_vector": query_vector}
            }
        }
    },
    "_source": {"includes": ["image_index", "category"]},
    "from": 0,
    "size": 10
}

res = es.search(index="raw-images", body=body)

{'took': 292,
 'timed_out': False,
 '_shards': {'total': 1, 'successful': 1, 'skipped': 0, 'failed': 0},
 'hits': {'total': {'value': 10000, 'relation': 'gte'},
  'max_score': 2.0,
  'hits': [{'_index': 'raw-images',
    '_type': 'images',
    '_id': '0rgjOW0B8Qw73ee7lR9P',
    '_score': 2.0,
    '_source': {'image_index': 0, 'category': 9}},
   {'_index': 'raw-images',
    '_type': 'images',
    '_id': 'SbgjOW0B8Qw73ee72oR1',
    '_score': 1.9564186,
    '_source': {'image_index': 25719, 'category': 9}},
   ....