0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

OpenSearch で 画像検索をする手順

Last updated at Posted at 2024-07-23

OpenSearch には k-NN インデックスをサポートしております。
また、k-NN インデックスに対してコサイン類似度を計算することができます。

そのため、下記のブログのように画像とテキストの Embedding を index に保存して、類似画像を検索することができます。

そこで、ここでは類似画像の検索を行う手順を記録しておきます。

前提

  • Amazon OpenSearch Service を用意しておきます。(ローカルOpenSearchやElasticSearchでもk-nn関連の機能がサポートされていたりプラグインがあれば動くと思います)
  • 今回は Sagemaker の notebook 環境で行いましたが、後述する GPU を使用しない方法でも可能だと思います
  • 各種必要なパッケージはインストールしてください。
  • (頻繁に機能改善が行われる分野でもあるので)間違いや修正した方が良い点があれば気軽にご連絡ください。

OpenSearch で類似画像検索の手順

  1. 今回は"cats_vs_dogs"データセットから画像データとラベルをダウンロードします

    from datasets import load_dataset
    
    # データセットをロード
    dataset = load_dataset("cats_vs_dogs")
    dataset
    
  2. OpenSearch の index に格納する画像データをデータセットから抽出します

    # 犬と猫の画像をそれぞれ10枚ずつ取得
    cat_images = [dataset["train"][idx]["image"] for idx in range(10) if dataset["train"][idx]["labels"] == 0]
    dog_images = [dataset["train"][idx]["image"] for idx in range(20000, 20010) if dataset["train"][idx]["labels"] == 1]
    

    dog_images[0]などで実際に画像データを確認できます。

  3. 画像から Embedding を抽出します

    import torchvision.transforms as transforms
    import torch
    import torchvision
    import timm
    
    # load model of EfficientNetb0
    model = timm.create_model('efficientnet_b0', pretrained=True) 
    model.classifier = torch.nn.Identity()
    model.eval()
    
    # preprocesser
    preprocess = transforms.Compose([
        transforms.Resize(256),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])
    
    def make_embed(image, preprocess=preprocess, model=model):
        inputs = preprocess(image)
        inputs = inputs.unsqueeze(0)
        embed = model(inputs)
        
        # 正規化
        embed = embed / torch.linalg.norm(embed, dim=1, keepdim=True)
        embed = embed[0, :] #(1, 1280) -> (1280)
        embed_list = embed.detach().cpu().numpy().tolist()
        return embed_list
    
    # create embedding and save embedding to list
    embed_dogs = []
    embed_cats = []
    for i in range(10):
        with torch.no_grad():
            dog_img = dog_images[i]
            cat_img = cat_images[i]
            
            dog_embed = make_embed(dog_img)
            cat_embed = make_embed(cat_img)
            
            embed_dogs.append(dog_embed)
            embed_cats.append(cat_embed)
    

    今回は Efficientnetb0 を使用します。
    model.classifier = torch.nn.Identity() とすることで、モデルのヘッドであるクラス分類層を削除して、embedding のみ取得することができます。

  4. OpenSearch クライアントに接続する

    from opensearchpy import OpenSearch, RequestsHttpConnection
    from requests_aws4auth import AWS4Auth
    from requests.auth import HTTPBasicAuth
    
    host = <host># e.g. my-test-domain.us-east-1.es.amazonaws.com
    awsauth=HTTPBasicAuth(<マスターユーザー名>, <マスターユーザーのパスワード>)
    
    # Create OpenSearch client
    client = OpenSearch(
        hosts = [{'host': host, 'port': 443}],
        http_auth = awsauth,
        use_ssl = True,
        verify_certs = True,
        connection_class = RequestsHttpConnection
    )
    
    # check to connect OpenSearch client
    client.info()
    

    今回は OpenSearch-py を OpenSearch クライアントとして使用しましたが、クライアントは何でも良いです。

  5. index 作成

    index_settings = {
      "settings": {
        "index.knn": True,
        "index.knn.space_type": "cosinesimil" # using cosine similarity
      },
      "mappings": {
        "properties": {
          "<任意の名前(今回は「embeddings_name」とします。)>": {
            "type": "knn_vector",
            "dimension": 1280
          }
        }
      }
    }
    
    # create index
    client.indices.create(
        index=<index 名>,
        body=index_settings
    )
    
  6. indexing

    # save animal images.
    corr_table = {}
    
    # using for corr_table's key
    ids = 1
    
    for i in range(10):
        dog_embed = embed_dogs[i]
        cat_embed = embed_cats[i]
        
        # body_data of dog
        dog_embedding_data = {
            'ids': ids,
            'label': 0,
            'embeddings_name':  dog_embed
        }
        corr_table[ids] = dog_images[i]
        ids += 1
    
        # body_data of cat
        cat_embedding_data = {
            'ids': ids,
            'label': 1,
            'embeddings_name': cat_embed
        }
        corr_table[ids] = cat_images[i]
        ids += 1
    
        # indexing
        client.index(
            index=<index 名>,
            body=dog_embedding_data,
            refresh=True
        )
        client.index(
            index=<index 名>,
            body=cat_embedding_data,
            refresh=True
        )
    
  7. 適当な画像データを用いてコサイン類似度で検索する

    # sampling image!
    input_image = dataset["train"][100]["image"]
    
    # create embedding
    with torch.no_grad():
        output_embed = make_embed(input_image)
    
    # search query body
    query = {
        "size": 20,
        "query": {
            "knn": {
                "embeddings_name": {
                    "vector": output_embed,
                    "k": 20,
                }
            }
        },
    }   
    
    # searching
    response = client.search(
        body=query,
        index=<index 名>,
    )
    
    # display the result of search
    for hit in response['hits']['hits']:
        score = hit['_score']
        doc = hit['_source']
        print(f"Score: {score}, ID: {doc['ids']}, Label: {doc['label']}")
    

    -> 以下のように検索されました。Label を確認すると、label==1である猫が上位に検索されており、類似度検索ができました!!!
    犬と猫は似ている気がするのでスコアが近いですが、きちんと猫の画像で検索すると猫が上位に来ており、正確に検索できてることが確認できました。
    Score: 0.68988866, ID: 8, Label: 1 Score: 0.6791921, ID: 6, Label: 1 Score: 0.6593978, ID: 2, Label: 1 Score: 0.65817195, ID: 16, Label: 1 Score: 0.65811276, ID: 10, Label: 1 Score: 0.6543478, ID: 18, Label: 1 Score: 0.6213816, ID: 20, Label: 1 Score: 0.61332446, ID: 12, Label: 1 Score: 0.60994273, ID: 14, Label: 1 Score: 0.58420175, ID: 4, Label: 1 Score: 0.5604887, ID: 17, Label: 0 Score: 0.5598484, ID: 13, Label: 0 Score: 0.5588222, ID: 19, Label: 0 Score: 0.5570838, ID: 3, Label: 0 Score: 0.54466915, ID: 1, Label: 0 Score: 0.53117865, ID: 15, Label: 0 Score: 0.51422846, ID: 5, Label: 0 Score: 0.50965375, ID: 9, Label: 0 Score: 0.5043833, ID: 7, Label: 0 Score: 0.5031408, ID: 11, Label: 0

  8. 実際に画像を比較します(input_imageとcorr_table[8]を比較してください)
    実際に類似してる猫の画像が検索できました!!!

まとめ

ここまでの手順で OpenSearch で類似画像検索ができました。
かなり簡単に検索するところまでいけました。例えば、テキストと画像のマルチモーダルでの類似してるものを検索することも簡単にできそうですね。

0
0
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?