Uso de la IA para encontrar a tu estilista de celebridades (Parte I)
El artículo fue publicado originalmente en el sitio web de AI Accelerator Institute y se vuelve a publicar aquí con permiso.
Me encanta llevar pantalones rosas durante mis charlas, lo que se ha convertido en uno de mis looks característicos. Tengo varios pares, lo que ha generado muchas conversaciones sobre moda y estilo. Recientemente, he participado en un proyecto llamado "Fashion AI", en el que utilizamos un modelo ajustado finamente para segmentar prendas de ropa en imágenes. Luego recortamos cada artículo etiquetado y redimensionamos las imágenes al mismo tamaño. Por último, almacenamos los embeddings generados a partir de esas imágenes en Milvus, una base de datos vectorial de código abierto que puede almacenar y consultar miles de millones de embeddings vectoriales.
Para encontrar los artículos con mayor coincidencia en nuestra base de datos, aplicamos las mismas transformaciones a la imagen y a la consulta a lo largo de los mismos vectores. Para cada consulta, este proyecto devuelve tres resultados. Puedes interpretar los resultados según tu preferencia. También puedes determinar qué celebridad es la coincidencia más cercana para ti. Puedes elegir el primer puesto más frecuente, la distancia agregada más baja o el más común en general.
Puedes encontrar las imágenes aquí. Además de las imágenes, necesitarás una versión actualizada de Python y pip install milvus pymilvus torch torchvision matplotlib. Usamos el modelo segmentador de ropa de Mateusz Dziemian en Hugging Face, y este modelo ResNet50 de Nvidia en PyTorch para la segmentación de imágenes y los embeddings.
En esta publicación, analizaremos cómo generar segmentación de imágenes para artículos de moda, añadir tus datos de imagen a Milvus y descubrir a qué celebridad se parece más tu vestido.
Segmentación de imágenes para prendas de ropa
Para realizar segmentación de imágenes, he encontrado tres modelos para examinar en Hugging Face.
El modelo Fashion-CLIP de Patrick John Chia
Finalmente elegí el modelo "segformer". Proporciona una segmentación precisa para diferentes prendas de ropa e identifica 18 tipos de "objetos". Por ejemplo, detecta "upper clothes" para cualquier tipo de partes superiores, "dress", "left shoe", "right shoe", "hat" y muchas más prendas de ropa. Además, puede detectar cosas como "face", "hair", "right leg" y "left leg". Puedes encontrar el conjunto completo de 18 tipos de objetos aquí.
Comenzamos importando los paquetes necesarios para la manipulación de imágenes en este proyecto. Estos incluyen torch para la extracción de características, el objeto segformer de transformers, matplotlib y algunas importaciones de torchvision como Resize, masks_to_boxes y 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
Generar máscaras de segmentación con Hugging Face
Hay muchos enfoques para segmentar tu imagen, dependiendo del modelo que uses y de lo que detecte. Para este ejemplo, nuestro modelo devuelve una imagen de 18 capas, una para cada tipo de objeto, incluido el fondo. La primera función que necesitamos escribir es una que genere esta imagen.
La función get_segmentation requiere tres parámetros: un extractor de características, un modelo y una imagen. Primero, genera características de entrada usando la imagen y el extractor. Luego, obtiene la salida del modelo y la convierte en logits. Después, sobremuestrea los logits mediante una interpolación bilineal de PyTorch. Finalmente, la función toma solo la predicción máxima para cada píxel en las muestras sobremuestreadas para crear una máscara de segmentación.
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
Como referencia, las imágenes en upsampled_logits se ven así:
Mientras que la imagen pred_seg se ve así: (estas son dos imágenes diferentes, aunque ambas son de Andre 3000).
Obtener las máscaras de segmentación es sencillo a partir de aquí. Obtenemos todos los valores únicos en la segmentación; en este modelo, solo puede haber hasta 18. Descartamos la primera entrada, que representa el fondo. Para crear las máscaras, extraemos los píxeles en la segmentación que tienen el mismo valor que el ID del objeto. Hago que esta función devuelva tanto las máscaras como los ID para que podamos hacer seguimiento de 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 función crea máscaras que se ven así (se muestran las máscaras de cabello y ropa superior):
Recorta y redimensiona tus imágenes con transformaciones de Pytorch
Ahora podemos crear una nueva imagen para cada objeto detectado usando las máscaras y las ideas de objeto de la función get_masks, así como la imagen original. Luego, llamamos a la mágica función masks_to_boxes, que importamos anteriormente desde torchvision.ops, para convertir las máscaras creadas en cuadros delimitadores.
A continuación, creamos una lista de cuadros para recortar y convertimos el sistema de coordenadas de los cuadros al sistema de coordenadas de crop. Los cuadros se devuelven como valores en la forma de (x1, x2, y1, y2). Por otro lado, la función crop espera una entrada en la forma de (top, left, height, width).
Antes de recortar las imágenes, también definimos una función de preprocesamiento. Queremos redimensionar cada imagen a 256x256 y convertirlas en tensores de PyTorch (actualmente son imágenes PIL). Ahora es momento de recortar las imágenes. Recorremos los cuadros de recorte y llamamos a la función crop sobre la imagen usando los valores que obtuvimos antes. Luego agregamos la imagen preprocesada como el valor correspondiente al valor clave del id de segmentación a un diccionario. Al final de la función, devolvemos ese diccionario.
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
A continuación se muestra un ejemplo de los cuadros que recortamos y para los que creamos imágenes separadas usando una salida de Drake in a fire.
Agrega tus datos de imagen a una base de datos vectorial
Ahora que tenemos todas las imágenes segmentadas y recortadas, agreguémoslas a Milvus, nuestra base de datos vectorial. Para ayudarte a comenzar rápidamente con Milvus, usamos Milvus Lite, una versión ligera de Milvus, en este ejemplo para ejecutar una instancia de Milvus en nuestro notebook. Luego usamos pymilvus para conectarnos al servidor predeterminado proporcionado por Milvus Lite.
También usamos esta sección para configurar algunas constantes. Definamos el número de dimensiones en un vector (del modelo Nvidia ResNet50), el tamaño del lote, el nombre de nuestra colección y el número de resultados a devolver. Finalmente, ejecutamos una función ssl para crear un contexto no verificado para obtener el modelo de 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
Definir tu esquema para almacenar metadatos en una base de datos vectorial
Paso uno: define tu esquema. El esquema se usa para organizar los datos guardados en la base de datos vectorial. El campo id es un ID de clave normal en bases de datos SQL o NoSQL, mientras que los otros campos tienen definiciones similares a SQL en sus tipos de datos (int64, varchar, float, etc).
Para este ejemplo, guardamos la ruta del archivo, el nombre de la celebridad y el ID de segmentación como metadatos. En el futuro, podríamos agregar más campos, como la ubicación de cuadros delimitadores o máscaras. Una vez que definimos el FieldSchema, definimos un CollectionSchema y luego creamos una Collection en Milvus basada en el esquema y el nombre de colección dados.
Ahora que tenemos una colección, definamos su índice. Estos parámetros de índice son bastante básicos. Usamos IVF Flat con 128 centroides y L2 como métrica de distancia. Creamos el índice en nuestra colección, especificando que el campo embedding es sobre el que se debe operar. Luego, cargamos la colección en memoria para que esté lista para operar con ella.
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()
Obtener tus embeddings vectoriales desde ResNet50 de Nvidia
El primer paso en esta sección es cargar el modelo. Cargamos el modelo ResNet50 de Nvidia desde PyTorch y luego eliminamos la capa de salida. Los embeddings vectoriales son la salida de la penúltima capa de un 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 función es responsable de recibir los embeddings vectoriales e insertar los datos en Milvus. Acepta tres parámetros: los datos, el objeto de colección y un modelo, que en este caso es el modelo de embeddings. Para hacer un seguimiento de cómo se está operando con los datos a medida que los agregamos a nuestra base de datos vectorial, he añadido varias sentencias print.
Además de imprimir datos de depuración, apilamos todos los valores en data[0] en un solo tensor y luego eliminamos cualquier dimensión de tamaño 1 de la salida usando la función squeeze. Luego, insertamos una nueva lista compuesta por las últimas tres entradas del lote de datos original, seguidas del tensor de salida convertido en una lista. Estas corresponden a la ruta del archivo, el nombre, el ID de segmentación y el embedding de 2048 dimensiones.
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()])
Los datos impresos se ven como la imagen que se muestra a continuación. Cada lote de datos tiene un tamaño de 128 hasta el final, y cada entrada tiene un tamaño de 3x256x256. La salida es un Tensor de PyTorch de longitud 128, y cada entrada en la salida tiene una longitud de 2048. El tensor impreso es la salida de la primera entrada en el lote de datos.
Almacenar tus datos de imagen en una base de datos vectorial
¿Recuerdas ese extractor y modelo de segmentación del que hablamos antes? Aquí es donde los usamos. Usamos este modelo segformer preentrenado de Hugging Face. Después de cargar los modelos, ponemos todas las rutas de archivo en una lista para iterar sobre ellas.
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}")
Milvus espera una lista de listas como entrada. En este ejemplo, usamos una lista de 4 listas, que corresponden a la imagen, la ruta del archivo, el nombre y el ID de segmentación. En la función embed_insert, convertimos la imagen en un embedding vectorial. Luego iteramos por cada ruta de archivo de las imágenes, recopilamos sus máscaras de segmentación y las recortamos. Finalmente, añadimos las imágenes con sus metadatos al lote de datos.
Cada 128 imágenes, las embebemos e insertamos en Milvus, y luego vaciamos el lote de datos. Al final del bucle, embebemos e insertamos el resto del lote de datos en Milvus y lo descargamos para completar la indexación. En un Mac M1 de 2021 con 16 GB de RAM, este proceso tarda aproximadamente 8 minutos en ejecutarse.
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()
Descubre a qué celebridades te vistes más parecido
Hay mucho que puedes hacer con esta configuración. Proporcionaré métodos adicionales para emparejar y evaluar tus elecciones de moda en un próximo artículo. En este ejemplo, obtendremos las tres mejores imágenes basadas en cada prenda segmentada. Usamos un par de ejemplos de Taylor Swift y obtenemos una recuperación perfecta.
Generar embeddings para tus imágenes de entrada
De forma similar a cómo cargamos imágenes en la base de datos, necesitamos procesar imágenes de entrada. La función para incrustar imágenes de búsqueda toma dos parámetros: datos y el modelo (de embeddings). Usamos el modelo para obtener los embeddings, los aplanamos o los comprimimos según el número de imágenes consultadas, los convertimos en una lista y los devolvemos.
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()
De forma similar a la función embed_insert, agregué varias instrucciones print aquí para hacer seguimiento de los datos. Como se muestra a continuación, los data pasados a esta función son esencialmente el objeto data[0], en comparación con la función embed_insert.
Para consultar la base de datos, solo necesitamos los embeddings vectoriales, que podemos obtener de una manera similar a cuando agregamos imágenes a Milvus. Sin embargo, es útil mantener estas otras variables en memoria para facilitar comparaciones más adelante.
# 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)
Consultar la base de datos vectorial
Ahora que tenemos los embeddings, podemos consultar la base de datos. Por diversión, estoy agregando el módulo time para registrar cuánto tardan estas consultas. En este ejemplo, medimos el tiempo de consulta para 23 vectores de 2048 dimensiones. Para consultar Milvus, simplemente usamos la función search con los embeddings generados arriba.
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)
Después de recorrer los resultados, podemos ver la respuesta generada, mostrada en la imagen debajo del código.
for index, result in enumerate(res):
print(index)
print(result)
Resumen
Eso es todo por esta sección. Ahora estás listo para comparar cualquier imagen tuya o de tus amigos (¡con permiso!) con algunas celebridades, incluidas Taylor Swift, Drake y Andre 3000. Para hacer esto, empieza obteniendo una segmentación de las prendas de vestir en la imagen usando un modelo encontrado en Hugging Face.
Con las segmentaciones en mano, toma cada segmentación única de la imagen y recórtalas en imágenes separadas. Antes de poner estas imágenes recortadas en la base de datos vectorial, redimensiónalas y conviértelas en tensores. Luego, pásalas por un modelo de embeddings ResNet50 de Nvidia para obtener los embeddings vectoriales que se almacenarán.
Para la consulta, realiza un procedimiento similar al de cargar los vectores. En este ejemplo, solo llegamos hasta obtener los resultados de la consulta. Para ir más allá, guarda los cuadros delimitadores o las máscaras en la base de datos vectorial y extráelos para mostrar coincidencias específicas. Como alternativa, vuelve a pasar las imágenes de entrada por el modelo y haz lo mismo. Como hacemos todo localmente, podemos usar la memoria local.
Espero que hayas disfrutado esto. No dudes en conectarte conmigo y compartir tus comentarios. Además, ¡por favor dime a qué celebridad te pareces más al vestirte!
¡Y no olvides consultar la Parte II de esta serie!
Sigue leyendo
Zilliz Cloud Now Available in AWS Asia Pacific (Seoul)
Zilliz Cloud is now available in AWS Seoul — low-latency vector search, in-country data residency, and one-step migration for Korean AI teams. 31 regions across 5 clouds.
Legal Document Analysis: Harnessing Zilliz Cloud's Semantic Search and RAG for Legal Insights
Enhance legal document analysis with Zilliz Cloud’s Semantic Search and RAG. Improve accuracy, efficiency, and scalability for contracts, case law, and compliance.
DeepRAG: Thinking to Retrieval Step by Step for Large Language Models
Discover DeepRAG, an advanced retrieval-augmented generation (RAG) model that improves LLM accuracy by retrieving only essential data through step-by-step reasoning.