Usare l'IA per trovare il tuo stylist delle celebrità (Parte I)
L'articolo è stato originariamente pubblicato sul sito web AI Accelerator Institute e viene ripubblicato qui con autorizzazione.
Adoro indossare pantaloni rosa durante i miei interventi, cosa che è diventata uno dei miei look distintivi. Ne possiedo diverse paia, il che ha dato il via a molte discussioni su moda e stile. Di recente, sono stato coinvolto in un progetto chiamato "Fashion AI," in cui utilizziamo un modello fine-tuned per segmentare i capi di abbigliamento nelle immagini. Poi ritagliamo ogni articolo etichettato e ridimensioniamo le immagini alla stessa dimensione. Infine, memorizziamo gli embedding generati da quelle immagini in Milvus, un database vettoriale open-source che può archiviare e interrogare miliardi di embedding vettoriali.
Per trovare gli articoli più simili nel nostro database, applichiamo le stesse trasformazioni all'immagine e alla query lungo gli stessi vettori. Per ogni query, questo progetto restituisce tre risultati. Puoi interpretare i risultati in base alle tue preferenze. Puoi anche determinare quale celebrità è l'abbinamento più vicino a te. Puoi scegliere il primo posto più comune, la distanza aggregata più bassa o il più comune complessivamente.
Puoi trovare le immagini qui. Oltre alle immagini, avrai bisogno di una versione aggiornata di Python e di pip install milvus pymilvus torch torchvision matplotlib. Usiamo il modello di segmentazione degli abiti di Mateusz Dziemian su Hugging Face, e questo modello ResNet50 di Nvidia su PyTorch per la segmentazione delle immagini e gli embedding.
In questo post, discuteremo come generare la segmentazione delle immagini per articoli di moda, aggiungere i dati delle tue immagini a Milvus e scoprire a quale celebrità assomiglia di più il tuo vestito.
Segmentazione delle immagini per capi di abbigliamento
Per eseguire la segmentazione delle immagini, ho trovato tre modelli da esaminare su Hugging Face.
Il modello Fashion-CLIP di Patrick John Chia
Alla fine ho scelto il modello "segformer". Fornisce una segmentazione accurata per diversi capi di abbigliamento e identifica 18 tipi di "oggetti." Ad esempio, rileva "upper clothes" per qualsiasi tipo di parte superiore, "dress," "left shoe," "right shoe," "hat," e molti altri capi di abbigliamento. Inoltre, può rilevare elementi come "face," "hair," "right leg," e "left leg." Puoi trovare l'insieme completo dei 18 tipi di oggetti qui.
Iniziamo importando i pacchetti necessari per la manipolazione delle immagini in questo progetto. Questi includono torch per l'estrazione delle caratteristiche, l'oggetto segformer da transformers, matplotlib, e alcuni import di torchvision come 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
Generare maschere di segmentazione con Hugging Face
Esistono molti approcci per segmentare la tua immagine, a seconda del modello che usi e di ciò che rileva. Per questo esempio, il nostro modello restituisce un'immagine a 18 livelli, uno per ciascun tipo di oggetto, incluso lo sfondo. La prima funzione che dobbiamo scrivere è quella che genera questa immagine.
La funzione get_segmentation richiede tre parametri: un estrattore di caratteristiche, un modello e un'immagine. Per prima cosa, genera le feature di input usando l'immagine e l'estrattore. Poi, ottiene l'output del modello e lo converte in logit. Successivamente, esegue l'upsampling dei logit tramite un'interpolazione bilineare di PyTorch. Infine, la funzione prende solo la previsione massima per ogni pixel nei campioni sottoposti a upsampling per creare una maschera di segmentazione.
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
Per riferimento, le immagini in upsampled_logits hanno questo aspetto:
Mentre l'immagine pred_seg ha questo aspetto: (si tratta di due immagini diverse, anche se entrambe sono di Andre 3000).
Ottenere le maschere di segmentazione da qui è semplice. Otteniamo tutti i valori univoci nella segmentazione; in questo modello, possono essercene al massimo 18. Scartiamo la prima voce, che rappresenta lo sfondo. Per creare le maschere, estraiamo i pixel nella segmentazione che hanno lo stesso valore dell'ID dell'oggetto. Faccio in modo che questa funzione restituisca sia le maschere sia gli ID, così possiamo tener traccia di entrambi.
# 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
Questa funzione crea maschere che hanno questo aspetto (sono mostrate le maschere dei capelli e degli indumenti superiori):
Ritaglia e ridimensiona le tue immagini con le trasformazioni Pytorch
Ora possiamo creare una nuova immagine per ogni oggetto rilevato usando le maschere e gli ID degli oggetti della funzione get_masks, oltre all'immagine originale. Quindi, chiamiamo la magica funzione masks_to_boxes, che abbiamo importato in precedenza da torchvision.ops, per convertire le maschere create in bounding box.
Successivamente, creiamo un elenco di box da ritagliare e convertiamo il sistema di coordinate dei box nel sistema di coordinate di crop. I box vengono restituiti come valori nella forma (x1, x2, y1, y2). D'altra parte, la funzione crop si aspetta un input nella forma (top, left, height, width).
Prima di ritagliare le immagini, definiamo anche una funzione di preprocessing. Vogliamo ridimensionare ogni immagine a 256x256 e convertirle in tensori PyTorch (attualmente immagini PIL). Ora è il momento di ritagliare le immagini. Eseguiamo un ciclo sui box di ritaglio e chiamiamo la funzione crop sull'immagine usando i valori ottenuti in precedenza. Poi aggiungiamo l'immagine preelaborata come valore corrispondente al valore chiave dell'id di segmentazione a un dizionario. Alla fine della funzione, restituiamo quel dizionario.
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
Di seguito è riportato un esempio dei box che ritagliamo e per cui creiamo immagini separate usando Drake in un output con fuoco.
Aggiungi i dati delle tue immagini a un database vettoriale
Ora che abbiamo tutte le immagini segmentate e ritagliate, aggiungiamole a Milvus, il nostro database vettoriale. Per aiutarti a iniziare rapidamente con Milvus, in questo esempio usiamo Milvus Lite, una versione leggera di Milvus, per eseguire un’istanza di Milvus nel nostro notebook. Poi usiamo pymilvus per connetterci al server predefinito fornito da Milvus Lite.
Usiamo anche questa sezione per configurare alcune costanti. Definiamo il numero di dimensioni in un vettore (dal modello Nvidia ResNet50), la dimensione del batch, il nome della nostra collection e il numero di risultati da restituire. Infine, eseguiamo una funzione ssl per creare un contesto non verificato per ottenere il modello da 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
Definire il tuo schema per archiviare i metadati in un database vettoriale
Primo passo: definire il tuo schema. Lo schema viene usato per organizzare i dati salvati nel database vettoriale. Il campo id è un normale ID chiave nei database SQL o NoSQL, mentre gli altri campi hanno definizioni simili a SQL nei loro tipi di dati (int64, varchar, float, ecc.).
Per questo esempio, salviamo il percorso del file, il nome della celebrità e l’ID di segmentazione come metadati. In futuro potremmo aggiungere altri campi, come la posizione dei bounding box o delle maschere. Una volta definito il FieldSchema, definiamo un CollectionSchema e poi creiamo una Collection in Milvus in base allo schema e al nome della collection forniti.
Ora che abbiamo una collection, definiamo il suo indice. Questi parametri dell’indice sono piuttosto basilari. Usiamo IVF Flat con 128 centroidi e L2 come metrica di distanza. Creiamo l’indice nella nostra collection, specificando che il campo embedding è quello su cui operare. Poi carichiamo la collection in memoria in modo che sia pronta per essere utilizzata.
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()
Ottenere i tuoi vector embedding da ResNet50 di Nvidia
Il primo passo in questa sezione è caricare il modello. Carichiamo il modello ResNet50 di Nvidia da PyTorch, poi rimuoviamo il layer di output. I vector embedding sono l’output del penultimo layer in un modello.
# 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()
Questa funzione è responsabile della ricezione dei vector embedding e dell’inserimento dei dati in Milvus. Accetta tre parametri: i dati, l’oggetto collection e un modello, che in questo caso è il modello di embedding. Per tenere traccia di come i dati vengono elaborati mentre li aggiungiamo al nostro database vettoriale, ho aggiunto diverse istruzioni print.
Oltre alla stampa dei dati di debug, impiliamo tutti i valori in data[0] in un unico tensore e poi rimuoviamo eventuali dimensioni di dimensione 1 dall’output usando la funzione squeeze. Quindi inseriamo una nuova lista composta dalle ultime tre voci del batch di dati originale, seguite dal tensore di output convertito in una lista. Questi corrispondono al percorso del file, al nome, all’ID di segmentazione e all’embedding a 2048 dimensioni.
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()])
I dati stampati appaiono come l’immagine mostrata di seguito. Ogni batch di dati ha una dimensione di 128 fino alla fine, con ogni voce di dimensione 3x256x256. L’output è un Tensor PyTorch di lunghezza 128, con ogni voce nell’output di lunghezza 2048. Il tensore stampato è l’output della prima voce nel batch di dati.
Archiviazione dei dati delle immagini in un database vettoriale
Ricordi il modello di estrazione e segmentazione di cui abbiamo parlato prima? È qui che li usiamo. Usiamo questo modello segformer pre-addestrato di Hugging Face. Dopo aver caricato i modelli, inseriamo tutti i percorsi dei file in una lista da scorrere.
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 si aspetta una lista di liste come input. In questo esempio, usiamo una lista di 4 liste, che corrispondono all’immagine, al percorso del file, al nome e all’ID di segmentazione. Nella funzione embed_insert, convertiamo l’immagine in un embedding vettoriale. Poi scorriamo ogni percorso di file verso le immagini, raccogliamo le loro maschere di segmentazione e le ritagliamo. Infine, aggiungiamo le immagini con i relativi metadati al batch di dati.
Ogni 128 immagini, le trasformiamo in embedding e le inseriamo in Milvus, quindi svuotiamo il batch di dati. Alla fine del ciclo, trasformiamo in embedding e inseriamo il resto del batch di dati in Milvus e lo svuotiamo per completare l’indicizzazione. Su un Mac M1 del 2021 con 16 GB di RAM, questo processo richiede circa 8 minuti per essere eseguito.
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()
Scopri a quali celebrità assomigli di più nel modo di vestire
C’è molto che puoi fare con questa configurazione. Fornirò metodi aggiuntivi per abbinare e valutare le tue scelte di moda in un prossimo articolo. In questo esempio, otterremo le prime tre immagini in base a ciascun capo di abbigliamento segmentato. Usiamo un paio di esempi di Taylor Swift e otteniamo un richiamo perfetto.
Generazione di embedding per le tue immagini di input
In modo simile a come carichiamo le immagini nel database, dobbiamo elaborare le immagini di input. La funzione per incorporare le immagini di ricerca accetta due parametri: data e il modello (di embedding). Usiamo il modello per ottenere gli embedding, li appiattiamo o li comprimiamo a seconda del numero di immagini interrogate, li convertiamo in una lista e li restituiamo.
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()
In modo simile alla funzione embed_insert, ho aggiunto qui diverse istruzioni print per tenere traccia dei dati. Come mostrato di seguito, i data passati a questa funzione sono essenzialmente l'oggetto data[0], rispetto alla funzione embed_insert.
Per interrogare il database, abbiamo bisogno solo degli embedding vettoriali, che possiamo ottenere in modo simile a quando abbiamo aggiunto immagini a Milvus. Tuttavia, è utile mantenere queste altre variabili in memoria per facilitare i confronti in seguito.
# 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)
Interrogare il database vettoriale
Ora che abbiamo gli embedding, possiamo interrogare il database. Per divertimento, aggiungo il modulo time per tenere traccia di quanto tempo impiegano queste query. In questo esempio, misuriamo il tempo di query per 23 vettori a 2048 dimensioni. Per interrogare Milvus, usiamo semplicemente la funzione search con gli embedding generati sopra.
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)
Dopo aver iterato sui risultati, possiamo vedere la risposta generata, mostrata nell'immagine sotto il codice.
for index, result in enumerate(res):
print(index)
print(result)
Riepilogo
Questo è tutto per questa sezione. Ora sei pronto a confrontare qualsiasi immagine di te stesso o dei tuoi amici (con permesso!) con alcune celebrità, tra cui Taylor Swift, Drake e Andre 3000. Per farlo, inizia ottenendo una segmentazione dei capi di abbigliamento nell'immagine usando un modello trovato su Hugging Face.
Con le segmentazioni a disposizione, prendi ogni segmentazione unica dall'immagine e ritagliale in immagini separate. Prima di inserire queste immagini ritagliate nel database vettoriale, ridimensionalle e trasformale in tensori. Quindi, passale attraverso un modello di embedding ResNet50 di Nvidia per ottenere gli embedding vettoriali da memorizzare.
Per eseguire query, esegui una procedura simile al caricamento dei vettori. In questo esempio, ci siamo spinti solo fino a ottenere i risultati della query. Per andare oltre, salva i riquadri di delimitazione o le maschere nel database vettoriale e recuperali per mostrare corrispondenze specifiche. In alternativa, fai passare di nuovo le immagini di input attraverso il modello e fai la stessa cosa. Poiché facciamo tutto localmente, possiamo usare la memoria locale.
Spero che ti sia piaciuto. Sentiti libero di connetterti con me e condividere il tuo feedback. Inoltre, fammi sapere a quale celebrità ti vesti più simile!
E non dimenticare di dare un'occhiata alla Parte II di questa serie!
Continua a leggere
Stop Building AI Data Infra for the Wrong Stage
Learn how AI data infrastructure should evolve from prototype to enterprise scale, and when Vector Lakebase becomes the right architecture for AI apps.
Build for the Boom: Why AI Agent Startups Should Build Scalable Infrastructure Early
Explore strategies for developing AI agents that can handle rapid growth. Don't let inadequate systems undermine your success during critical breakthrough moments.
VidTok: Rethinking Video Processing with Compact Tokenization
VidTok tokenizes videos to reduce redundancy while preserving spatial and temporal details for efficient processing.