Pesquisa de imagens com o Milvus
Neste bloco de notas, vamos mostrar-lhe como utilizar o Milvus para procurar imagens semelhantes num conjunto de dados. Para o demonstrar, utilizaremos um subconjunto do conjunto de dados ImageNet e procuraremos uma imagem de um cão afegão.
Preparação do conjunto de dados
Primeiro, precisamos de carregar o conjunto de dados e extraí-lo para processamento posterior.
!wget https://github.com/milvus-io/pymilvus-assets/releases/download/imagedata/reverse_image_search.zip
!unzip -q -o reverse_image_search.zip
Pré-requisitos
Para executar este notebook, é necessário ter as seguintes dependências instaladas:
- pymilvus>=2.4.2
- timm
- torch
- numpy
- sklearn
- almofada
Para executar o Colab, fornecemos os comandos práticos para instalar as dependências necessárias.
$ pip install pymilvus --upgrade
$ pip install timm
Se estiver a utilizar o Google Colab, para ativar as dependências acabadas de instalar, poderá ter de reiniciar o tempo de execução. (Clique no menu "Runtime" (Tempo de execução) na parte superior do ecrã e selecione "Restart session" (Reiniciar sessão) no menu pendente).
Definir o Extrator de caraterísticas
Em seguida, precisamos definir um extrator de caraterísticas que extrai a incorporação de uma imagem usando o modelo ResNet-34 do timm.
import torch
from PIL import Image
import timm
from sklearn.preprocessing import normalize
from timm.data import resolve_data_config
from timm.data.transforms_factory import create_transform
class FeatureExtractor:
def __init__(self, modelname):
# Load the pre-trained model
self.model = timm.create_model(
modelname, pretrained=True, num_classes=0, global_pool="avg"
)
self.model.eval()
# Get the input size required by the model
self.input_size = self.model.default_cfg["input_size"]
config = resolve_data_config({}, model=modelname)
# Get the preprocessing function provided by TIMM for the model
self.preprocess = create_transform(**config)
def __call__(self, imagepath):
# Preprocess the input image
input_image = Image.open(imagepath).convert("RGB") # Convert to RGB if needed
input_image = self.preprocess(input_image)
# Convert the image to a PyTorch tensor and add a batch dimension
input_tensor = input_image.unsqueeze(0)
# Perform inference
with torch.no_grad():
output = self.model(input_tensor)
# Extract the feature vector
feature_vector = output.squeeze().numpy()
return normalize(feature_vector.reshape(1, -1), norm="l2").flatten()
Criar uma coleção Milvus
Em seguida, temos de criar uma coleção Milvus para armazenar os embeddings da imagem
from pymilvus import MilvusClient
# Set up a Milvus client
client = MilvusClient(uri="example.db")
# Create a collection in quick setup mode
if client.has_collection(collection_name="image_embeddings"):
client.drop_collection(collection_name="image_embeddings")
client.create_collection(
collection_name="image_embeddings",
vector_field_name="vector",
dimension=512,
auto_id=True,
enable_dynamic_field=True,
metric_type="COSINE",
)
Quanto ao argumento de MilvusClient
:
- Definir o
uri
como um ficheiro local, por exemplo,./milvus.db
, é o método mais conveniente, uma vez que utiliza automaticamente o Milvus Lite para armazenar todos os dados neste ficheiro. - Se tiver uma grande escala de dados, pode configurar um servidor Milvus mais eficiente em docker ou kubernetes. Nesta configuração, utilize o uri do servidor, por exemplo,
http://localhost:19530
, como o seuuri
. - Se pretender utilizar o Zilliz Cloud, o serviço de nuvem totalmente gerido para o Milvus, ajuste os endereços
uri
etoken
, que correspondem ao Public Endpoint e à chave Api no Zilliz Cloud.
Inserir os embeddings no Milvus
Iremos extrair os embeddings de cada imagem utilizando o modelo ResNet34 e inserir as imagens do conjunto de treino no Milvus.
import os
extractor = FeatureExtractor("resnet34")
root = "./train"
insert = True
if insert is True:
for dirpath, foldername, filenames in os.walk(root):
for filename in filenames:
if filename.endswith(".JPEG"):
filepath = dirpath + "/" + filename
image_embedding = extractor(filepath)
client.insert(
"image_embeddings",
{"vector": image_embedding, "filename": filepath},
)
from IPython.display import display
query_image = "./test/Afghan_hound/n02088094_4261.JPEG"
results = client.search(
"image_embeddings",
data=[extractor(query_image)],
output_fields=["filename"],
search_params={"metric_type": "COSINE"},
)
images = []
for result in results:
for hit in result[:10]:
filename = hit["entity"]["filename"]
img = Image.open(filename)
img = img.resize((150, 150))
images.append(img)
width = 150 * 5
height = 150 * 2
concatenated_image = Image.new("RGB", (width, height))
for idx, img in enumerate(images):
x = idx % 5
y = idx // 5
concatenated_image.paste(img, (x * 150, y * 150))
display("query")
display(Image.open(query_image).resize((150, 150)))
display("results")
display(concatenated_image)
'query'
png
'results'
Resultados
Podemos ver que a maioria das imagens são da mesma categoria que a imagem de pesquisa, que é o cão de caça afegão. Isto significa que encontrámos imagens semelhantes à imagem de pesquisa.
Implementação rápida
Para saber como iniciar uma demonstração online com este tutorial, consulte o aplicativo de exemplo.