Usando IA para encontrar seu stylist de celebridades (Parte I)
O artigo foi originalmente publicado no site do AI Accelerator Institute e é republicado aqui com permissão.
Adoro usar calças rosa durante minhas palestras, o que se tornou uma das minhas marcas registradas. Tenho vários pares, o que gerou muitas discussões sobre moda e estilo. Recentemente, estive envolvido em um projeto chamado "Fashion AI", no qual utilizamos um modelo ajustado para segmentar roupas em imagens. Em seguida, recortamos cada peça rotulada e redimensionamos as imagens para o mesmo tamanho. Por fim, armazenamos os embeddings gerados a partir dessas imagens no Milvus, um banco de dados vetorial de código aberto que pode armazenar e consultar bilhões de embeddings vetoriais.
Para encontrar as peças mais parecidas em nosso banco de dados, aplicamos as mesmas transformações à imagem e à consulta ao longo dos mesmos vetores. Para cada consulta, este projeto retorna três resultados. Você pode interpretar os resultados com base na sua preferência. Você também pode determinar qual celebridade é a combinação mais próxima de você. Você pode escolher o primeiro lugar mais comum, a menor distância agregada ou o mais comum no geral.
Você pode encontrar as imagens aqui. Além das imagens, você precisará de uma versão atualizada do Python e pip install milvus pymilvus torch torchvision matplotlib. Usamos o modelo segmentador de roupas de Mateusz Dziemian no Hugging Face, e este modelo ResNet50 da Nvidia no PyTorch para segmentação de imagens e embeddings.
Neste post, discutiremos como gerar segmentação de imagens para itens de moda, adicionar seus dados de imagem ao Milvus e descobrir com qual celebridade seu vestido mais se parece.
Segmentação de imagens para peças de roupa
Para realizar a segmentação de imagens, encontrei três modelos para analisar no Hugging Face.
O modelo Fashion-CLIP de Patrick John Chia
No fim, escolhi o modelo "segformer". Ele fornece segmentação precisa para diferentes peças de roupa e identifica 18 tipos de "objetos." Por exemplo, ele detecta "upper clothes" para qualquer tipo de parte de cima, "dress," "left shoe," "right shoe," "hat," e muitas outras peças de roupa. Além disso, ele pode detectar coisas como "face," "hair," "right leg," e "left leg." Você pode encontrar o conjunto completo de 18 tipos de objeto aqui.
Começamos importando os pacotes necessários para manipulação de imagens neste projeto. Eles incluem torch para extração de recursos, o objeto segformer de transformers, matplotlib, e algumas importações de torchvision, como Resize, masks_to_boxes, e crop.
import torch
from torch import nn, tensor
from transformers import AutoFeatureExtractor, SegformerForSemanticSegmentation
import matplotlib.pyplot as plt
from torchvision.transforms import Resize
import torchvision.transforms as T
from torchvision.ops import masks_to_boxes
from torchvision.transforms.functional import crop
Gerar máscaras de segmentação com Hugging Face
Há muitas abordagens para segmentar sua imagem, dependendo do modelo que você usa e do que ele detecta. Para este exemplo, nosso modelo retorna uma imagem de 18 camadas, uma para cada tipo de objeto, incluindo o plano de fundo. A primeira função que precisamos escrever é uma que gere essa imagem.
A função get_segmentation requer três parâmetros: um extrator de características, um modelo e uma imagem. Primeiro, ela gera características de entrada usando a imagem e o extrator. Em seguida, obtém a saída do modelo e a converte em logits. Depois, faz o upsample dos logits por meio de uma interpolação bilinear do PyTorch. Por fim, a função usa apenas a previsão máxima para cada pixel nas amostras com upsample para criar uma máscara de segmentação.
def get_segmentation(extractor, model, image):
inputs = extractor(images=image, return_tensors="pt")
outputs = model(**inputs)
logits = outputs.logits.cpu()
upsampled_logits = nn.functional.interpolate(
logits,
size=image.size[::-1],
mode="bilinear",
align_corners=False,
)
pred_seg = upsampled_logits.argmax(dim=1)[0]
return pred_seg
Para referência, as imagens em upsampled_logits se parecem com isto:
Enquanto a imagem pred_seg se parece com isto: (estas são duas imagens diferentes, embora ambas sejam de Andre 3000).
Obter as máscaras de segmentação é simples a partir daqui. Obtemos todos os valores únicos na segmentação; neste modelo, só pode haver até 18. Descartamos a primeira entrada, que representa o plano de fundo. Para criar as máscaras, extraímos os pixels na segmentação que têm o mesmo valor que o ID do objeto. Faço esta função retornar tanto as máscaras quanto os IDs para que possamos acompanhar ambos.
# returns two lists masks (tensor) and obj_ids (int)
# "mattmdjaga/segformer_b2_clothes" from hugging face
def get_masks(segmentation):
obj_ids = torch.unique(segmentation)
obj_ids = obj_ids[1:]
masks = segmentation == obj_ids[:, None, None]
return masks, obj_ids
Esta função cria máscaras que se parecem com isto (máscaras de cabelo e roupas superiores mostradas):
Recorte e redimensione suas imagens com transforms do Pytorch
Agora podemos criar uma nova imagem para cada objeto detectado usando as máscaras e os IDs de objeto da função get_masks, bem como a imagem original. Em seguida, chamamos a função mágica masks_to_boxes, que importamos anteriormente de torchvision.ops, para converter as máscaras criadas em caixas delimitadoras.
Depois, criamos uma lista de caixas para recortar e convertemos o sistema de coordenadas das caixas para o sistema de coordenadas de crop. As caixas são retornadas como valores no formato (x1, x2, y1, y2). Por outro lado, a função crop espera uma entrada no formato (top, left, height, width).
Antes de recortar as imagens, também definimos uma função de pré-processamento. Queremos redimensionar cada imagem para 256x256 e convertê-las em Tensores do PyTorch (atualmente Imagens PIL). Agora é hora de recortar as imagens. Percorremos as caixas de recorte e chamamos a função crop na imagem usando os valores que obtivemos anteriormente. Em seguida, adicionamos a imagem pré-processada como o valor correspondente ao valor-chave do ID de segmentação a um dicionário. No final da função, retornamos esse dicionário.
def crop_images(masks, obj_ids, img):
boxes = masks_to_boxes(masks)
crop_boxes = []
for box in boxes:
crop_box = tensor([box[0], box[1], box[2]-box[0], box[3]-box[1]])
crop_boxes.append(crop_box)
preprocess = T.Compose([
T.Resize(size=(256, 256)),
T.ToTensor()
])
cropped_images = {}
for i in range(len(crop_boxes)):
crop_box = crop_boxes[i]
cropped = crop(img, crop_box[1].item(), crop_box[0].item(), crop_box[3].item(), crop_box[2].item())
cropped_images[obj_ids[i].item()] = preprocess(cropped)
return cropped_images
Abaixo está um exemplo das caixas que recortamos e para as quais criamos imagens separadas usando Drake em uma saída com fogo.
Adicione seus dados de imagem a um banco de dados vetorial
Agora que temos todas as imagens segmentadas e recortadas, vamos adicioná-las ao Milvus, nosso banco de dados vetorial. Para ajudar você a começar rapidamente com o Milvus, usamos o Milvus Lite, uma versão leve do Milvus, neste exemplo para executar uma instância do Milvus em nosso notebook. Em seguida, usamos pymilvus para nos conectar ao servidor padrão fornecido pelo Milvus Lite.
Também usamos esta seção para configurar algumas constantes. Vamos definir o número de dimensões em um vetor (do modelo Nvidia ResNet50), o tamanho do lote, o nome da nossa coleção e o número de resultados a retornar. Por fim, executamos uma função ssl para criar um contexto não verificado para obter o modelo do PyTorch.
from milvus import default_server
from pymilvus import utility, connections
default_server.start()
connections.connect(host="127.0.0.1", port=default_server.listen_port)
DIMENSION = 2048
BATCH_SIZE = 128
COLLECTION_NAME = "fashion"
TOP_K = 3
# run this before importing th resnet50 model if you run into an SSL certificate URLError
import ssl
ssl._create_default_https_context = ssl._create_unverified_context
Definindo seu esquema para armazenar metadados em um banco de dados vetorial
Primeiro passo: defina seu esquema. O esquema é usado para organizar os dados salvos no banco de dados vetorial. O campo id é um ID de chave comum em bancos de dados SQL ou NoSQL, enquanto os outros campos têm definições semelhantes às de SQL em seus tipos de dados (int64, varchar, float, etc).
Neste exemplo, salvamos o caminho do arquivo, o nome da celebridade e o ID de segmentação como metadados. No futuro, podemos adicionar mais campos, como a localização de caixas delimitadoras ou máscaras. Depois de definirmos o FieldSchema, definimos um CollectionSchema e, em seguida, criamos uma Collection no Milvus com base no esquema e no nome da coleção fornecidos.
Agora que temos uma coleção, vamos definir seu índice. Esses parâmetros de índice são bem básicos. Usamos IVF Flat com 128 centróides e L2 como a métrica de distância. Criamos o índice em nossa coleção, especificando que o campo embedding é aquele sobre o qual operar. Em seguida, carregamos a coleção na memória para que ela esteja pronta para ser operada.
from pymilvus import FieldSchema, CollectionSchema, Collection, DataType
fields = [
FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
FieldSchema(name='filepath', dtype=DataType.VARCHAR, max_length=200),
FieldSchema(name="name", dtype=DataType.VARCHAR, max_length=200),
FieldSchema(name="seg_id", dtype=DataType.INT64),
FieldSchema(name='embedding', dtype=DataType.FLOAT_VECTOR, dim=DIMENSION)
]
schema = CollectionSchema(fields=fields)
collection = Collection(name=COLLECTION_NAME, schema=schema)
index_params = {
"index_type": "IVF_FLAT",
"metric_type": "L2",
"params": {"nlist": 128},
}
collection.create_index(field_name="embedding", index_params=index_params)
collection.load()
Obtendo seus embeddings vetoriais do ResNet50 da Nvidia
O primeiro passo nesta seção é carregar o modelo. Carregamos o modelo ResNet50 da Nvidia a partir do PyTorch e, em seguida, removemos a camada de saída. Embeddings vetoriais são a saída da penúltima camada em um modelo.
# Load the embedding model with the last layer removed
embeddings_model = torch.hub.load('NVIDIA/DeepLearningExamples:torchhub', 'nvidia_resnet50', pretrained=True)
embeddings_model = torch.nn.Sequential(*(list(embeddings_model.children())[:-1]))
embeddings_model.eval()
Esta função é responsável por receber os embeddings vetoriais e inserir os dados no Milvus. Ela aceita três parâmetros: os dados, o objeto de coleção e um modelo, que neste caso é o modelo de embedding. Para acompanhar como os dados estão sendo operados à medida que os adicionamos ao nosso banco de dados vetorial, adicionei várias instruções print.
Além de imprimir dados de depuração, empilhamos todos os valores em data[0] em um tensor e, em seguida, removemos quaisquer dimensões de tamanho 1 da saída usando a função squeeze. Depois, inserimos uma nova lista composta pelas três últimas entradas do lote de dados original, seguidas pelo tensor de saída convertido em uma lista. Elas correspondem ao caminho do arquivo, nome, ID de segmentação e ao embedding de 2048 dimensões.
def embed_insert(data, collection, model):
with torch.no_grad():
print(len(data[0]))
print(data[0][0].size())
output = model(torch.stack(data[0])).squeeze()
print(type(output))
print(len(output))
print(len(output[0]))
print(output[0])
collection.insert([data[1], data[2], data[3], output.tolist()])
Os dados impressos se parecem com a imagem mostrada abaixo. Cada lote de dados tem um tamanho de 128 até o final, com cada entrada sendo de tamanho 3x256x256. A saída é um Tensor do PyTorch de comprimento 128, com cada entrada na saída tendo comprimento 2048. O tensor impresso é a saída da primeira entrada no lote de dados.
Armazenando seus dados de imagem em um banco de dados vetorial
Lembra daquele modelo extrator e de segmentação sobre o qual falamos antes? É aqui que os usamos. Usamos este modelo segformer pré-treinado da Hugging Face. Depois de carregar os modelos, colocamos todos os caminhos dos arquivos em uma lista para percorrê-los.
extractor = AutoFeatureExtractor.from_pretrained("mattmdjaga/segformer_b2_clothes")
model = SegformerForSemanticSegmentation.from_pretrained("mattmdjaga/segformer_b2_clothes")
import os
image_paths = []
for celeb in os.listdir("./photos"):
for image in os.listdir(f"./photos/{celeb}/"):
# print(image)
image_paths.append(f"./photos/{celeb}/{image}")
O Milvus espera uma lista de listas como entrada. Neste exemplo, usamos uma lista de 4 listas, que correspondem à imagem, ao caminho do arquivo, ao nome e ao ID de segmentação. Na função embed_insert, convertemos a imagem em um embedding vetorial. Em seguida, percorremos cada caminho de arquivo para as imagens, coletamos suas máscaras de segmentação e as recortamos. Por fim, adicionamos as imagens com seus metadados ao lote de dados.
A cada 128 imagens, geramos embeddings e as inserimos no Milvus, e então limpamos o lote de dados. No final do loop, geramos embeddings e inserimos o restante do lote de dados no Milvus e fazemos flush para concluir a indexação. Em um Mac M1 2021 com 16 GB de RAM, esse processo leva aproximadamente 8 minutos para executar.
from PIL import Image
data_batch = [[], [], [], []]
for path in image_paths:
image = Image.open(path)
path_split = path.split("/")
name = " ".join(path_split[2].split("_"))
segmentation = get_segmentation(extractor, model, image)
masks, ids = get_masks(segmentation)
cropped_images = crop_images(masks, ids, image)
for key, image in cropped_images.items():
data_batch[0].append(image)
data_batch[1].append(path)
data_batch[2].append(name)
data_batch[3].append(key)
if len(data_batch[0]) % BATCH_SIZE == 0:
embed_insert(data_batch, collection, embeddings_model)
data_batch = [[], [], [], []]
if len(data_batch[0]) != 0:
embed_insert(data_batch, collection, embeddings_model)
collection.flush()
Descubra com quais celebridades você se veste de forma mais parecida
Há muito que você pode fazer com esta configuração. Fornecerei métodos adicionais para combinar e avaliar suas escolhas de moda em um próximo texto. Neste exemplo, obteremos as três principais fotos com base em cada peça de roupa segmentada. Usamos alguns exemplos de Taylor Swift e obtemos uma revocação perfeita.
Gerando embeddings para suas imagens de entrada
Semelhante à forma como carregamos imagens no banco de dados, precisamos processar imagens de entrada. A função para incorporar imagens de busca recebe dois parâmetros: data e o modelo (de embedding). Usamos o modelo para obter os embeddings, achatá-los ou espremê-los dependendo do número de imagens consultadas, convertê-los em uma lista e retorná-los.
def embed_search_images(data, model):
with torch.no_grad():
print(len(data[0]))
print(data[0][0].size())
output = model(torch.stack(data))
print(type(output))
print(len(output))
print(len(output[0]))
print(output[0])
if len(output) > 1:
return output.squeeze().tolist()
Else:
return torch.flatten(output, start_dim=1).tolist()
Semelhante à função embed_insert, adicionei várias instruções print aqui para acompanhar os dados. Como mostrado abaixo, os data passados para esta função são essencialmente o objeto data[0], em comparação com a função embed_insert.
Para consultar o banco de dados, precisamos apenas dos embeddings vetoriais, que podemos obter de maneira semelhante a quando adicionamos imagens ao Milvus. No entanto, é útil manter essas outras variáveis na memória para facilitar comparações posteriormente.
# data_batch[0] is a list of tensors
# data_batch[1] is a list of filepaths to the images (string)
# data_batch[2] is a list of the names of the people in the images (string)
# data_batch[3] is a list of segmentation keys (int)
data_batch = [[], [], [], []]
search_paths = ["./photos/Taylor_Swift/Taylor_Swift_3.jpg", "./photos/Taylor_Swift/Taylor_Swift_8.jpg"]
for path in search_paths:
image = Image.open(path)
path_split = path.split("/")
name = " ".join(path_split[2].split("_"))
segmentation = get_segmentation(extractor, model, image)
masks, ids = get_masks(segmentation)
cropped_images = crop_images(masks, ids, image)
for key, image in cropped_images.items():
data_batch[0].append(image)
data_batch[1].append(path)
data_batch[2].append(name)
data_batch[3].append(key)
embeds = embed_search_images(data_batch[0], embeddings_model)
Consultando o banco de dados vetorial
Agora que temos os embeddings, podemos consultar o banco de dados. Por diversão, estou adicionando o módulo time para acompanhar quanto tempo essas consultas levam. Neste exemplo, medimos o tempo de consulta para 23 vetores de 2048 dimensões. Para consultar o Milvus, simplesmente usamos a função search com os embeddings gerados acima.
import time
start = time.time()
res = collection.search(embeds,
anns_field='embedding',
param={"metric_type": "L2",
"params": {"nprobe": 10}},
limit=TOP_K,
output_fields=['filepath'])
finish = time.time()
print(finish - start)
Depois de percorrer os resultados, podemos ver a resposta gerada, mostrada na imagem abaixo do código.
for index, result in enumerate(res):
print(index)
print(result)
Resumo
Isso é tudo para esta seção. Agora você está pronto para comparar quaisquer imagens suas ou de seus amigos (com permissão!) com algumas celebridades, incluindo Taylor Swift, Drake e Andre 3000. Para fazer isso, comece obtendo uma segmentação das peças de roupa na imagem usando um modelo encontrado no Hugging Face.
Com as segmentações em mãos, pegue cada segmentação única da imagem e recorte-as em imagens separadas. Antes de colocar essas imagens recortadas no banco de dados vetorial, redimensione-as e transforme-as em tensores. Em seguida, passe-as por um modelo de embeddings ResNet50 da Nvidia para obter os embeddings vetoriais a serem armazenados.
Para consultas, realize um procedimento semelhante ao carregamento dos vetores. Neste exemplo, fomos apenas até obter os resultados da consulta. Para ir além, salve as caixas delimitadoras ou máscaras no banco de dados vetorial e extraia-as para mostrar correspondências específicas. Como alternativa, execute as imagens de entrada pelo modelo novamente e faça a mesma coisa. Como fazemos tudo localmente, podemos usar a memória local.
Espero que você tenha gostado disso. Sinta-se à vontade para conectar-se comigo e compartilhar seu feedback. Além disso, por favor, me diga com qual celebridade você se veste mais parecido!
E não se esqueça de conferir a Parte II desta série!
Continue lendo
How to Choose the Best Embedding Model for RAG in 2026: 10 Models Benchmarked
We benchmarked 10 embedding models on cross-modal, cross-lingual, long-document, and dimension compression tasks. See which one fits your RAG pipeline.
Announcing VDBBench 1.0: Open-Source VectorDB Benchmarking with Your Real-World Production Workloads
Discover VDBBench 1.0, an open-source tool for benchmarking vector databases with real-world production data, streaming ingestion, and concurrent workloads.
Vector Databases vs. Time Series Databases
Use a vector database for similarity search and semantic relationships; use a time series database for tracking value changes over time.