Pesquisa de imagens com o Milvus

Nesta página, vamos analisar um exemplo simples de pesquisa de imagens usando Milvus. O conjunto de dados que estamos a pesquisar é o Impressionist-Classifier Dataset encontrado no Kaggle. Para este exemplo, voltámos a alojar os dados numa unidade pública do Google.

Para este exemplo, estamos apenas a utilizar o modelo Resnet50 pré-treinado da Torchvision para os embeddings. Vamos começar!

Instalar os requisitos

Para este exemplo, vamos utilizar pymilvus para nos ligarmos ao Milvus, torch para executar o modelo de incorporação, torchvision para o modelo atual e o pré-processamento, gdown para descarregar o conjunto de dados de exemplo e tqdm para carregar barras.

pip install pymilvus torch gdown torchvision tqdm

Recolher os dados

Vamos utilizar gdown para obter o zip do Google Drive e depois descomprimi-lo com a biblioteca integrada 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")

Argumentos globais

Estes são alguns dos principais argumentos globais que iremos utilizar para facilitar o acompanhamento e a atualização.

# 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

Configurando o Milvus

Neste ponto, vamos começar a configurar o Milvus. Os passos são os seguintes:

1. Ligue-se à instância do Milvus utilizando o URI fornecido.

   from pymilvus import connections
   
   # Connect to the instance
   connections.connect(host=MILVUS_HOST, port=MILVUS_PORT)
   

2. Se a coleção já existir, elimine-a.

   from pymilvus import utility
   
   # Remove any previous collections with the same name
   if utility.has_collection(COLLECTION_NAME):
       utility.drop_collection(COLLECTION_NAME)
   

3. Crie a coleção que contém o ID, o caminho do ficheiro da imagem e a sua incorporação.

   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. Crie um índice na coleção recém-criada e carregue-a para a memória.

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

Uma vez concluídos estes passos, a coleção está pronta para ser inserida e pesquisada. Todos os dados adicionados serão indexados automaticamente e ficarão imediatamente disponíveis para pesquisa. Se os dados forem muito recentes, a pesquisa pode ser mais lenta, uma vez que a pesquisa por força bruta será utilizada em dados que ainda estão a ser indexados.

Inserir os dados

Para este exemplo, vamos utilizar o modelo ResNet50 fornecido por torch e o seu hub de modelos. Para obter os embeddings, estamos a retirar a camada de classificação final, o que faz com que o modelo nos dê embeddings de 2048 dimensões. Todos os modelos de visão encontrados em torch utilizam o mesmo pré-processamento que incluímos aqui.

Nos próximos passos, vamos fazer o seguinte:

1. Carregar os dados.

   import glob
   
   # Get the filepaths of the images
   paths = glob.glob('./paintings/paintings/**/*.jpg', recursive=True)
   len(paths)
   

2. Pré-processamento dos dados em lotes.

   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. Incorporar os dados.

   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. Inserir os dados.

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

   • Esta etapa é relativamente demorada porque a incorporação leva tempo. Tome um gole de café e relaxe.
   • O PyTorch pode não funcionar bem com o Python 3.9 e versões anteriores. Considere usar o Python 3.10 e versões posteriores.

Efetuar a pesquisa

Com todos os dados inseridos no Milvus, podemos começar a efetuar as nossas pesquisas. Neste exemplo, vamos procurar duas imagens de exemplo. Como estamos a fazer uma pesquisa em lote, o tempo de pesquisa é partilhado entre as imagens do lote.

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

A imagem do resultado da pesquisa deve ser semelhante à seguinte:

Resultado da pesquisa de imagens