Similarité d’image avec les ensembles de données et les transformateurs Hugging Face

Recherche de similarité d'image avec Hugging Face Transformers et des ensembles de données.

Dans ce post, vous apprendrez à construire un système de similarité d’images avec 🤗 Transformers. Découvrir la similarité entre une image de requête et des candidats potentiels est un cas d’utilisation important pour les systèmes de recherche d’informations, tels que la recherche d’images inversée, par exemple. Tout ce que le système essaie de répondre, c’est que, étant donné une image de requête et un ensemble d’images candidates, quelles images sont les plus similaires à l’image de requête.

Nous utiliserons la bibliothèque 🤗 datasets car elle prend en charge de manière transparente le traitement parallèle, ce qui sera utile lors de la construction de ce système.

Bien que le post utilise un modèle basé sur ViT ( nateraw/vit-base-beans ) et un ensemble de données particulier ( Beans ), il peut être étendu à l’utilisation d’autres modèles prenant en charge la modalité vision et d’autres ensembles de données d’images. Voici quelques modèles notables que vous pourriez essayer:

• Bienvenue à PaddlePaddle sur le Hugging Face Hub
• Segmentation universelle des images avec Mask2Former et OneFormer
• Génération d’actifs 3D IA pour le développement de jeux #3

• Swin Transformer
• ConvNeXT
• RegNet

De plus, l’approche présentée dans le post peut potentiellement être étendue à d’autres modalités également.

Pour étudier le système de similarité d’images entièrement fonctionnel, vous pouvez vous référer au cahier Colab lié au début.

Comment définissons-nous la similarité?

Pour construire ce système, nous devons d’abord définir comment nous voulons calculer la similarité entre deux images. Une pratique largement répandue consiste à calculer des représentations denses (embeddings) des images données, puis à utiliser la similarité cosinus pour déterminer à quel point les deux images sont similaires.

Pour ce post, nous utiliserons les “embeddings” pour représenter les images dans l’espace vectoriel. Cela nous donne une manière agréable de compresser de manière significative l’espace de pixels en haute dimension des images (par exemple 224 x 224 x 3) en quelque chose de beaucoup plus basse dimension (768, par exemple). L’avantage principal de cela est le temps de calcul réduit dans les étapes suivantes.

Calcul des embeddings

Pour calculer les embeddings à partir des images, nous utiliserons un modèle de vision qui a une certaine compréhension de la façon de représenter les images d’entrée dans l’espace vectoriel. Ce type de modèle est également couramment appelé encodeur d’images.

Pour charger le modèle, nous utilisons la classe AutoModel. Elle fournit une interface pour charger n’importe quel point de contrôle de modèle compatible à partir du hub Hugging Face. En plus du modèle, nous chargeons également le processeur associé au modèle pour le prétraitement des données.

from transformers import AutoFeatureExtractor, AutoModel


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

Dans ce cas, le point de contrôle a été obtenu en affinant un modèle basé sur Vision Transformer sur l’ensemble de données beans.

Voici quelques questions qui pourraient se poser ici:

Q1 : Pourquoi n’avons-nous pas utilisé AutoModelForImageClassification ?

C’est parce que nous voulons obtenir des représentations denses des images et non des catégories discrètes, ce que AutoModelForImageClassification aurait fourni.

Q2 : Pourquoi ce point de contrôle en particulier?

Comme mentionné précédemment, nous utilisons un ensemble de données spécifique pour construire le système. Donc, au lieu d’utiliser un modèle généraliste (comme ceux entraînés sur l’ensemble de données ImageNet-1k, par exemple), il est préférable d’utiliser un modèle qui a été affiné sur l’ensemble de données utilisé. De cette façon, le modèle sous-jacent comprend mieux les images d’entrée.

Notez que vous pouvez également utiliser un point de contrôle qui a été obtenu grâce à un pré-entraînement auto-supervisé. Le point de contrôle ne doit pas nécessairement provenir de l’apprentissage supervisé. En fait, si elles sont bien pré-entraînées, les modèles auto-supervisés peuvent donner d’impressionnantes performances de recherche.

Maintenant que nous avons un modèle pour calculer les embeddings, nous avons besoin de certaines images candidates pour effectuer des requêtes.

Chargement d’un ensemble de données pour les images candidates

Dans un certain temps, nous allons construire des tables de hachage qui mappent les images candidates aux hachages. Pendant le temps de requête, nous utiliserons ces tables de hachage. Nous parlerons plus des tables de hachage dans la section respective, mais pour l’instant, pour avoir un ensemble d’images candidates, nous utiliserons la division train de l’ensemble de données beans.

from datasets import load_dataset


dataset = load_dataset("beans")

Voici à quoi ressemble un seul échantillon de la partie d’entraînement :

L’ensemble de données a trois caractéristiques :

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)}

Pour démontrer le système de similarité d’images, nous utiliserons 100 échantillons de l’ensemble de données d’images candidates pour réduire la durée d’exécution globale.

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

Le processus de recherche d’images similaires

Ci-dessous, vous pouvez trouver une vue d’ensemble picturale du processus sous-jacent à la recherche d’images similaires.

En décomposant un peu la figure ci-dessus, nous avons :

1. Extraire les embeddings des images candidates (candidate_subset), en les stockant dans une matrice.
2. Prendre une image de requête et extraire ses embeddings.
3. Parcourir la matrice d’embeddings (calculée à l’étape 1) et calculer le score de similarité entre l’embedding de requête et les embeddings candidats actuels. Nous maintenons généralement une correspondance sous forme de dictionnaire entre un identifiant de l’image candidate et les scores de similarité.
4. Trier la structure de correspondance par rapport aux scores de similarité et renvoyer les identifiants sous-jacents. Nous utilisons ces identifiants pour récupérer les échantillons candidats.

Nous pouvons écrire une simple fonction utilitaire et l’appliquer à notre ensemble de données d’images candidates pour calculer les embeddings de manière efficace.

import torch 

def extract_embeddings(model: torch.nn.Module):
    """Utilitaire pour calculer les embeddings."""
    device = model.device

    def pp(batch):
        images = batch["image"]
        # `transformation_chain` est une composition de transformations de prétraitement
        # que nous appliquons aux images d'entrée pour les préparer
        # pour le modèle. Pour plus de détails, consultez le cahier Colab accompagnant.
        image_batch_transformed = torch.stack(
            [transformation_chain(image) pour image dans 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

Et nous pouvons appliquer extract_embeddings() comme ceci :

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)

Ensuite, pour plus de commodité, nous créons une liste contenant les identifiants des images candidates.

candidate_ids = []

pour id dans tqdm(range(len(candidate_subset_emb))):
    label = candidate_subset_emb[id]["labels"]

    # Créer un identifiant unique.
    entry = str(id) + "_" + str(label)

    candidate_ids.append(entry)

Nous utiliserons la matrice des embeddings de toutes les images candidates pour calculer les scores de similarité avec une image de requête. Nous avons déjà calculé les embeddings des images candidates. Dans la cellule suivante, nous les rassemblons simplement dans une matrice.

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

Nous utiliserons la similarité cosinus pour calculer le score de similarité entre deux vecteurs d’embedding. Nous l’utiliserons ensuite pour récupérer des échantillons candidats similaires en fonction d’un échantillon de requête.

def compute_scores(emb_one, emb_two):
    """Calcule la similarité cosinus entre deux vecteurs."""
    scores = torch.nn.functional.cosine_similarity(emb_one, emb_two)
    return scores.numpy().tolist()


def fetch_similar(image, top_k=5):
    """Récupère les `top_k` images similaires avec `image` comme requête."""
    # Prépare l'image de requête en entrée pour le calcul de l'embedding.
    image_transformed = transformation_chain(image).unsqueeze(0)
    new_batch = {"pixel_values": image_transformed.to(device)}

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

    # Calcul des scores de similarité avec toutes les images candidates en une seule fois.
    # Nous créons également une correspondance entre les identifiants d'images candidates
    # et leurs scores de similarité avec l'image de requête.
    sim_scores = compute_scores(all_candidate_embeddings, query_embeddings)
    similarity_mapping = dict(zip(candidate_ids, sim_scores))
 
    # Trier le dictionnaire de correspondance et renvoyer les `top_k` candidats.
    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

Effectuer une requête

Avec toutes les fonctionnalités disponibles, nous sommes équipés pour effectuer une recherche de similarité. Prenons une image de requête à partir de la division test de l’ensemble de données 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"Label de la requête : {test_label}")
print(f"5 meilleures étiquettes candidates : {sim_labels}")

Cela donne :

Label de la requête : 0
5 meilleures étiquettes candidates : [0, 0, 0, 0, 0]

Il semble que notre système ait obtenu le bon ensemble d’images similaires. Lorsqu’elles sont visualisées, nous obtenons :

Extensions supplémentaires et conclusions

Nous disposons maintenant d’un système de similarité d’images fonctionnel. Mais en réalité, vous serez confronté à beaucoup plus d’images candidates. En tenant compte de cela, notre procédure actuelle présente plusieurs inconvénients :

• Si nous stockons les embeddings tels quels, les exigences en mémoire peuvent augmenter rapidement, en particulier lorsqu’il s’agit de millions d’images candidates. Les embeddings sont de dimension 768 dans notre cas, ce qui peut être relativement élevé dans le régime à grande échelle.
• Avoir des embeddings de haute dimension a un effet direct sur les calculs ultérieurs impliqués dans la partie de récupération.

Si nous pouvons réduire la dimensionalité des embeddings sans perturber leur signification, nous pouvons encore maintenir un bon compromis entre la vitesse et la qualité de la récupération. Le cahier Colab accompagnant cette publication met en œuvre et démontre des utilitaires pour y parvenir avec la projection aléatoire et le hachage localement sensible.

🤗 Datasets offre des intégrations directes avec FAISS, ce qui simplifie davantage le processus de construction de systèmes de similarité. Supposons que vous ayez déjà extrait les embeddings des images candidates (l’ensemble de données beans) et les ayez stockés dans une fonctionnalité appelée embeddings. Vous pouvez maintenant utiliser facilement la fonction add_faiss_index() de l’ensemble de données pour construire un index dense :

dataset_with_embeddings.add_faiss_index(column="embeddings")

Une fois l’index construit, dataset_with_embeddings peut être utilisé pour récupérer les exemples les plus proches en fonction des embeddings de requête avec get_nearest_examples() :

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

La méthode renvoie les scores et les exemples candidats correspondants. Pour en savoir plus, vous pouvez consulter la documentation officielle et ce cahier.

Enfin, vous pouvez essayer l’espace suivant qui construit une application miniaturisée de similarité d’images :

Dans cette publication, nous avons rapidement parcouru un guide de démarrage pour la construction de systèmes de similarité d’images. Si vous avez trouvé cette publication intéressante, nous vous recommandons vivement de vous appuyer sur les concepts que nous avons abordés ici afin de vous familiariser davantage avec leur fonctionnement interne.

Vous souhaitez en savoir plus ? Voici quelques ressources supplémentaires qui pourraient vous être utiles :

• Faiss : une bibliothèque pour une recherche de similarité efficace
• ScaNN : recherche de similarité vectorielle efficace
• Intégration de recherche d’images dans des applications mobiles

We will continue to update IPGirl; if you have any questions or suggestions, please contact us!