使用 Milvus 进行图像搜索

在本页中，我们将使用 Milvus 演示一个简单的图像搜索示例。我们搜索的数据集是Kaggle 上的印象派分类器数据集。在本示例中，我们将数据重新托管在谷歌公共驱动器中。

在本例中，我们只使用 Torchvision 预先训练好的 Resnet50 模型进行嵌入。让我们开始吧

安装要求

在本例中，我们将使用pymilvus 连接使用 Milvus，torch 运行嵌入模型，torchvision 实际模型和预处理，gdown 下载示例数据集，tqdm 加载条形图。

pip install pymilvus torch gdown torchvision tqdm

抓取数据

我们将使用gdown 从 Google Drive 抓取压缩包，然后使用内置的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")

全局参数

这些是我们将使用的一些主要全局参数，以便于跟踪和更新。

# 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. 如果 Collections 已经存在，则将其删除。

   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、图片文件路径及其 Embeddings 的 Collection。

   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. 在新创建的 Collections 上创建索引，并将其加载到内存中。

   # 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()
   

完成这些步骤后，就可以插入并搜索 Collections 了。任何添加的数据都会自动编入索引，并立即可供搜索。如果数据非常新，搜索速度可能会慢一些，因为将对仍在编制索引过程中的数据使用暴力搜索。

插入数据

在本例中，我们将使用torch 及其模型中心提供的 ResNet50 模型。为了获得 Embeddings，我们要去掉最后的分类层，这样模型就能为我们提供 2048 维的 embeddings。在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()
   

   • 这一步相对耗时，因为 Embeddings 需要时间。喝一口咖啡，放松一下。
   • PyTorch 可能无法在 Python 3.9 及更早版本中很好地运行。请考虑使用 Python 3.10 及更高版本。

执行搜索

将所有数据插入 Milvus 后，我们就可以开始执行搜索了。在本例中，我们将搜索两张示例图片。由于我们进行的是批量搜索，因此搜索时间由批量中的图像共享。

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')

搜索结果图像应与下图类似：

图像搜索输出