milvus-logo
LFAI
フロントページへ
  • 統合

Milvusを使った画像検索

このページでは、Milvusを使った簡単な画像検索の例を説明します。検索するデータセットはKaggleにあるImpressionist-Classifier Datasetです。この例では、データをgoogleドライブに再ホストしています。

この例では、埋め込みにTorchvisionで事前に訓練されたResnet50モデルを使用します。さっそく始めましょう!

要件のインストール

この例では、Milvusへの接続にpymilvus 、エンベッディングモデルの実行にtorch 、実際のモデルと前処理にtorchvision 、サンプルデータセットのダウンロードにgdown 、バーのロードにtqdm

pip install pymilvus torch gdown torchvision tqdm

データの取得

gdown を使ってGoogle Driveからzipを取得し、組み込みのzipfile ライブラリで解凍する。

import gdown
import zipfile

url = 'https://drive.google.com/uc?id=1OYDHLEy992qu5C4C8HV5uDIkOWRTAR1_'
output = './paintings.zip'
gdown.download(url, output)

with zipfile.ZipFile("./paintings.zip","r") as zip_ref:
    zip_ref.extractall("./paintings")

データセットのサイズは2.35GBで、ダウンロードにかかる時間はネットワーク状況に依存する。

グローバル引数

トラッキングとアップデートを容易にするために使用する主なグローバル引数である。

# Milvus Setup Arguments
COLLECTION_NAME = 'image_search'  # Collection name
DIMENSION = 2048  # Embedding vector size in this example
MILVUS_HOST = "localhost"
MILVUS_PORT = "19530"

# Inference Arguments
BATCH_SIZE = 128
TOP_K = 3

Milvusのセットアップ

この時点で、Milvusのセットアップを開始します。手順は以下の通りです:

  1. 提供されたURIを使用してMilvusインスタンスに接続する。

    from pymilvus import connections
    
    # Connect to the instance
    connections.connect(host=MILVUS_HOST, port=MILVUS_PORT)
    
  2. コレクションが既に存在する場合は、それを削除する。

    from pymilvus import utility
    
    # Remove any previous collections with the same name
    if utility.has_collection(COLLECTION_NAME):
        utility.drop_collection(COLLECTION_NAME)
    
  3. ID、画像のファイルパス、埋め込みを保持するコレクションを作成する。

    from pymilvus import FieldSchema, CollectionSchema, DataType, Collection
    
    # Create collection which includes the id, filepath of the image, and image embedding
    fields = [
        FieldSchema(name='id', dtype=DataType.INT64, is_primary=True, auto_id=True),
        FieldSchema(name='filepath', dtype=DataType.VARCHAR, max_length=200),  # VARCHARS need a maximum length, so for this example they are set to 200 characters
        FieldSchema(name='image_embedding', dtype=DataType.FLOAT_VECTOR, dim=DIMENSION)
    ]
    schema = CollectionSchema(fields=fields)
    collection = Collection(name=COLLECTION_NAME, schema=schema)
    
  4. 新しく作成したコレクションにインデックスを作成し、メモリにロードする。

    # Create an AutoIndex index for collection
    index_params = {
    'metric_type':'L2',
    'index_type':"IVF_FLAT",
    'params':{'nlist': 16384}
    }
    collection.create_index(field_name="image_embedding", index_params=index_params)
    collection.load()
    

これらの手順が完了すると、コレクションに挿入して検索する準備が整います。追加されたデータは自動的にインデックスが作成され、すぐに検索できるようになります。データが非常に新しい場合、まだインデックスが作成されていないデータに対して総当たり検索が使用されるため、検索が遅くなる可能性があります。

データの挿入

この例では、torch とそのモデルハブが提供する ResNet50 モデルを使用する。エンベッディングを得るために、最後の分類レイヤーを削除します。その結果、モデルは2048次元のエンベッディングを得ることになります。torch にあるビジョン・モデルはすべて、ここで紹介したのと同じ前処理を使用しています。

次のいくつかのステップでは

  1. データをロードする。

    import glob
    
    # Get the filepaths of the images
    paths = glob.glob('./paintings/paintings/**/*.jpg', recursive=True)
    len(paths)
    
  2. データをバッチに前処理する。

    import torch
    
    # Load the embedding model with the last layer removed
    model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet50', pretrained=True)
    model = torch.nn.Sequential(*(list(model.children())[:-1]))
    model.eval()
    
  3. データの埋め込み

    from torchvision import transforms
    
    # Preprocessing for images
    preprocess = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])
    
  4. データを挿入する。

    from PIL import Image
    from tqdm import tqdm
    
    # Embed function that embeds the batch and inserts it
    def embed(data):
        with torch.no_grad():
            output = model(torch.stack(data[0])).squeeze()
            collection.insert([data[1], output.tolist()])
    
    data_batch = [[],[]]
    
    # Read the images into batches for embedding and insertion
    for path in tqdm(paths):
        im = Image.open(path).convert('RGB')
        data_batch[0].append(preprocess(im))
        data_batch[1].append(path)
        if len(data_batch[0]) % BATCH_SIZE == 0:
            embed(data_batch)
            data_batch = [[],[]]
    
    # Embed and insert the remainder
    if len(data_batch[0]) != 0:
        embed(data_batch)
    
    # Call a flush to index any unsealed segments.
    collection.flush()
    
    • 埋め込みには時間がかかるので、このステップは比較的時間がかかる。コーヒーを一口飲んでリラックスしてください。
    • PyTorch は Python 3.9 以前のバージョンではうまく動かないかもしれません。代わりに Python 3.10 以降を使うことを検討してください。

すべてのデータがMilvusに挿入されたので、検索を開始することができます。この例では、2つの画像を検索します。バッチ検索を行っているため、検索時間はバッチの画像で共有されます。

import glob

# Get the filepaths of the search images
search_paths = glob.glob('./paintings/test_paintings/**/*.jpg', recursive=True)
len(search_paths)
import time
from matplotlib import pyplot as plt

# Embed the search images
def embed(data):
    with torch.no_grad():
        ret = model(torch.stack(data))
        # If more than one image, use squeeze
        if len(ret) > 1:
            return ret.squeeze().tolist()
        # Squeeze would remove batch for single image, so using flatten
        else:
            return torch.flatten(ret, start_dim=1).tolist()

data_batch = [[],[]]

for path in search_paths:
    im = Image.open(path).convert('RGB')
    data_batch[0].append(preprocess(im))
    data_batch[1].append(path)

embeds = embed(data_batch[0])
start = time.time()
res = collection.search(embeds, anns_field='image_embedding', param={'nprobe': 128}, limit=TOP_K, output_fields=['filepath'])
finish = time.time()
# Show the image results
f, axarr = plt.subplots(len(data_batch[1]), TOP_K + 1, figsize=(20, 10), squeeze=False)

for hits_i, hits in enumerate(res):
    axarr[hits_i][0].imshow(Image.open(data_batch[1][hits_i]))
    axarr[hits_i][0].set_axis_off()
    axarr[hits_i][0].set_title('Search Time: ' + str(finish - start))
    for hit_i, hit in enumerate(hits):
        axarr[hits_i][hit_i + 1].imshow(Image.open(hit.entity.get('filepath')))
        axarr[hits_i][hit_i + 1].set_axis_off()
        axarr[hits_i][hit_i + 1].set_title('Distance: ' + str(hit.distance))

# Save the search result in a separate image file alongside your script.
plt.savefig('search_result.png')

検索結果画像は以下のようになるはずです:

Image search output 画像検索の出力

翻訳DeepL

Try Managed Milvus for Free

Zilliz Cloud is hassle-free, powered by Milvus and 10x faster.

Get Started
フィードバック

このページは役に立ちましたか ?