Similarità delle immagini con Hugging Face Datasets e Transformers

Similarità immagini con Hugging Face Datasets e Transformers.

In questo post, imparerai a costruire un sistema di similarità delle immagini con 🤗 Transformers. Trovare la similarità tra un’immagine di query e dei potenziali candidati è un caso d’uso importante per i sistemi di recupero delle informazioni, come ad esempio la ricerca inversa delle immagini. Tutto ciò che il sistema cerca di rispondere è quale immagine di query e quale insieme di immagini candidato sono più simili all’immagine di query.

Sfrutteremo la libreria 🤗 datasets in quanto supporta in modo trasparente l’elaborazione parallela, che sarà utile nella costruzione di questo sistema.

Sebbene il post utilizzi un modello basato su ViT (nateraw/vit-base-beans) e un dataset specifico (Beans), può essere esteso per utilizzare altri modelli che supportano la modalità visione e altri dataset di immagini. Alcuni modelli notevoli che potresti provare:

  • Swin Transformer
  • ConvNeXT
  • RegNet

Inoltre, l’approccio presentato nel post può potenzialmente essere esteso ad altre modalità.

Per studiare il sistema di similarità delle immagini completamente funzionante, puoi fare riferimento al Notebook Colab collegato all’inizio.

Come definiamo la similarità?

Per costruire questo sistema, dobbiamo prima definire come vogliamo calcolare la similarità tra due immagini. Una pratica molto diffusa è calcolare rappresentazioni dense (embedding) delle immagini fornite e quindi utilizzare la metrica di similarità coseno per determinare quanto sono simili le due immagini.

In questo post, useremo “embedding” per rappresentare le immagini nello spazio vettoriale. Questo ci offre un modo efficace per comprimere in modo significativo lo spazio dei pixel ad alta dimensionalità delle immagini (ad esempio, 224 x 224 x 3) in qualcosa di molto più bassa dimensionalità (ad esempio, 768). Il vantaggio principale di ciò è il tempo di calcolo ridotto nei passaggi successivi.

Calcolo degli embedding

Per calcolare gli embedding dalle immagini, utilizzeremo un modello di visione che ha una certa comprensione di come rappresentare le immagini di input nello spazio vettoriale. Questo tipo di modello è comunemente chiamato anche codificatore di immagini.

Per caricare il modello, sfruttiamo la classe AutoModel. Essa fornisce un’interfaccia per caricare qualsiasi checkpoint di modello compatibile da Hugging Face Hub. Insieme al modello, carichiamo anche il processore associato al modello per la pre-elaborazione dei dati.

from transformers import AutoFeatureExtractor, AutoModel


model_ckpt = "nateraw/vit-base-beans"
extractor = AutoFeatureExtractor.from_pretrained(model_ckpt)
model = AutoModel.from_pretrained(model_ckpt)

In questo caso, il checkpoint è stato ottenuto tramite fine-tuning di un modello basato su Vision Transformer sul dataset beans.

Alcune domande che potrebbero sorgere qui:

Q1 : Perché non abbiamo usato AutoModelForImageClassification ?

Questo perché vogliamo ottenere rappresentazioni dense delle immagini e non categorie discrete, che è ciò che AutoModelForImageClassification avrebbe fornito.

Q2 : Perché proprio questo checkpoint?

Come accennato in precedenza, stiamo utilizzando un dataset specifico per costruire il sistema. Quindi, anziché utilizzare un modello generico (come quelli addestrati sul dataset ImageNet-1k, ad esempio), è meglio utilizzare un modello che è stato sottoposto a fine-tuning sul dataset utilizzato. In questo modo, il modello sottostante comprende meglio le immagini di input.

Nota che è anche possibile utilizzare un checkpoint ottenuto tramite pre-training auto-supervisionato. Il checkpoint non deve necessariamente provenire dall’apprendimento supervisionato. Infatti, se addestrati adeguatamente, i modelli auto-supervisionati possono fornire prestazioni di recupero impressionanti.

Ora che abbiamo un modello per calcolare gli embedding, abbiamo bisogno di alcune immagini candidato su cui effettuare le query.

Caricamento di un dataset per le immagini candidato

Tra un po’, costruiremo tabelle hash che mappano le immagini candidato agli hash. Durante il tempo di query, utilizzeremo queste tabelle hash. Parleremo più in dettaglio delle tabelle hash nella rispettiva sezione, ma per ora, per avere un insieme di immagini candidato, utilizzeremo la divisione train del dataset beans.

from datasets import load_dataset


dataset = load_dataset("beans")

Ecco come appare un singolo campione dal training split:

Il dataset ha tre caratteristiche:

dataset["train"].features
>>> {'image_file_path': Value(dtype='string', id=None),
 'image': Image(decode=True, id=None),
 'labels': ClassLabel(names=['angular_leaf_spot', 'bean_rust', 'healthy'], id=None)}

Per dimostrare il sistema di similarità delle immagini, utilizzeremo 100 campioni del dataset delle immagini candidate per mantenere breve il tempo di esecuzione complessivo.

num_samples = 100
seed = 42
candidate_subset = dataset["train"].shuffle(seed=seed).select(range(num_samples))

Il processo di ricerca di immagini simili

Di seguito, puoi trovare una panoramica grafica del processo che sta alla base del recupero di immagini simili.

Analizzando un po’ la figura precedente, abbiamo:

  1. Estrarre i vettori di embedding dalle immagini candidate (candidate_subset), memorizzandoli in una matrice.
  2. Prendere un’immagine di query ed estrarne il vettore di embedding.
  3. Iterare sulla matrice di embedding (calcolata al passaggio 1) e calcolare il punteggio di similarità tra l’embedding della query e gli embedding delle immagini candidate correnti. Solitamente manteniamo una mappatura simile a un dizionario che mantiene una corrispondenza tra un qualche identificatore dell’immagine candidata e i punteggi di similarità.
  4. Ordinare la struttura di mappatura rispetto ai punteggi di similarità e restituire gli identificatori sottostanti. Utilizziamo questi identificatori per recuperare i campioni candidati.

Possiamo scrivere una semplice utility e mapparla sul nostro dataset di immagini candidate per calcolare gli embedding in modo efficiente.

import torch 

def extract_embeddings(model: torch.nn.Module):
    """Utility per calcolare gli embedding."""
    device = model.device

    def pp(batch):
        images = batch["image"]
        # `transformation_chain` è una composizione di trasformazioni di preelaborazione
        # che applichiamo alle immagini di input per prepararle per il modello. Per ulteriori dettagli, consulta il Colab Notebook allegato.
        image_batch_transformed = torch.stack(
            [transformation_chain(image) for image in images]
        )
        new_batch = {"pixel_values": image_batch_transformed.to(device)}
        with torch.no_grad():
            embeddings = model(**new_batch).last_hidden_state[:, 0].cpu()
        return {"embeddings": embeddings}

    return pp

E possiamo mappare extract_embeddings() in questo modo:

device = "cuda" if torch.cuda.is_available() else "cpu"
extract_fn = extract_embeddings(model.to(device))
candidate_subset_emb = candidate_subset.map(extract_fn, batched=True, batch_size=batch_size)

Successivamente, per comodità, creiamo una lista contenente gli identificatori delle immagini candidate.

candidate_ids = []

for id in tqdm(range(len(candidate_subset_emb))):
    label = candidate_subset_emb[id]["labels"]

    # Creare un identificatore unico.
    entry = str(id) + "_" + str(label)

    candidate_ids.append(entry)

Utilizzeremo la matrice degli embedding di tutte le immagini candidate per calcolare i punteggi di similarità con un’immagine di query. Abbiamo già calcolato gli embedding delle immagini candidate. Nella cella successiva, li raccogliamo tutti insieme in una matrice.

all_candidate_embeddings = np.array(candidate_subset_emb["embeddings"])
all_candidate_embeddings = torch.from_numpy(all_candidate_embeddings)

Utilizzeremo la similarità coseno per calcolare il punteggio di similarità tra due vettori di embedding. Lo utilizzeremo quindi per recuperare campioni candidati simili dato un campione di query.

def compute_scores(emb_one, emb_two):
    """Calcola la similarità coseno tra due vettori."""
    scores = torch.nn.functional.cosine_similarity(emb_one, emb_two)
    return scores.numpy().tolist()


def fetch_similar(image, top_k=5):
    """Recupera le `top_k` immagini simili con `image` come query."""
    # Prepara l'immagine di query di input per il calcolo dell'embedding.
    image_transformed = transformation_chain(image).unsqueeze(0)
    new_batch = {"pixel_values": image_transformed.to(device)}

    # Calcola l'embedding.
    with torch.no_grad():
        query_embeddings = model(**new_batch).last_hidden_state[:, 0].cpu()

    # Calcola i punteggi di similarità con tutte le immagini candidate in un'unica volta.
    # Creiamo anche una mappatura tra gli identificatori delle immagini candidate
    # e i loro punteggi di similarità con l'immagine di query.
    sim_scores = compute_scores(all_candidate_embeddings, query_embeddings)
    similarity_mapping = dict(zip(candidate_ids, sim_scores))
 
    # Ordina il dizionario di mappatura e restituisce i `top_k` candidati.
    similarity_mapping_sorted = dict(
        sorted(similarity_mapping.items(), key=lambda x: x[1], reverse=True)
    )
    id_entries = list(similarity_mapping_sorted.keys())[:top_k]

    ids = list(map(lambda x: int(x.split("_")[0]), id_entries))
    labels = list(map(lambda x: int(x.split("_")[-1]), id_entries))
    return ids, labels

Esegui una query

Dato tutti gli strumenti, siamo attrezzati per fare una ricerca di similarità. Prendiamo un’immagine di query dal set di test del dataset “beans”:

test_idx = np.random.choice(len(dataset["test"]))
test_sample = dataset["test"][test_idx]["image"]
test_label = dataset["test"][test_idx]["labels"]

sim_ids, sim_labels = fetch_similar(test_sample)
print(f"Etichetta della query: {test_label}")
print(f"Principali 5 etichette candidate: {sim_labels}")

Conduce a:

Etichetta della query: 0
Principali 5 etichette candidate: [0, 0, 0, 0, 0]

Sembra che il nostro sistema abbia ottenuto il giusto insieme di immagini simili. Quando visualizzato, otterremmo:

Ulteriori estensioni e conclusioni

Abbiamo ora un sistema di similarità delle immagini funzionante. Ma nella realtà, si avranno molte più immagini candidate. Tenendo conto di ciò, la nostra procedura attuale ha diverse limitazioni:

  • Se memorizziamo gli embedding così come sono, i requisiti di memoria possono aumentare rapidamente, specialmente quando si lavora con milioni di immagini candidate. Gli embedding sono di dimensione 768-d nel nostro caso, che può ancora essere relativamente elevata nel regime di larga scala.
  • Avere embedding ad alta dimensionalità ha un effetto diretto sui calcoli successivi coinvolti nella parte di recupero.

Se riusciamo in qualche modo a ridurre la dimensionalità degli embedding senza alterarne il significato, possiamo comunque mantenere un buon compromesso tra velocità e qualità di recupero. Il Colab Notebook allegato a questo post implementa e dimostra strumenti per raggiungere questo obiettivo tramite proiezione casuale e locality-sensitive hashing.

🤗 Datasets offre integrazioni dirette con FAISS che semplificano ulteriormente il processo di costruzione di sistemi di similarità. Diciamo che hai già estratto gli embedding delle immagini candidate (il dataset “beans”) e li hai memorizzati in una caratteristica chiamata “embeddings”. Ora puoi facilmente utilizzare la funzione “add_faiss_index()” del dataset per costruire un indice denso:

dataset_with_embeddings.add_faiss_index(column="embeddings")

Una volta costruito l’indice, “dataset_with_embeddings” può essere utilizzato per recuperare gli esempi più vicini dati gli embedding di query con la funzione “get_nearest_examples()”:

scores, retrieved_examples = dataset_with_embeddings.get_nearest_examples(
    "embeddings", qi_embedding, k=top_k
)

Il metodo restituisce i punteggi e gli esempi candidati corrispondenti. Per saperne di più, puoi consultare la documentazione ufficiale e questo notebook.

Infine, puoi provare lo spazio seguente che costruisce un’applicazione mini di similarità delle immagini:

In questo post, abbiamo eseguito una panoramica per la costruzione di sistemi di similarità delle immagini. Se hai trovato interessante questo post, ti consigliamo vivamente di costruire sopra i concetti che abbiamo discusso qui in modo da poterti sentire più a tuo agio con il funzionamento interno.

Stai ancora cercando di imparare di più? Ecco alcune risorse aggiuntive che potrebbero esserti utili:

  • Faiss: una libreria per la ricerca di similarità efficiente
  • ScaNN: Ricerca efficiente di similarità vettoriale
  • Integrazione di cercatori di immagini all’interno delle applicazioni mobili