Construire un projet de correspondance d'images réel avec Gemini Embedding 2

Construire un projet de correspondance d'images réel avec Gemini Embedding 2

Gemini Embedding 2 : Caractéristiques clés

Les systèmes d'embedding traditionnels sont souvent conçus uniquement pour le texte. Si vous souhaitiez construire un système fonctionnant avec des images, de l'audio ou des documents, vous deviez généralement assembler plusieurs pipelines. Gemini Embedding 2 change cela en mappant différents types de contenu dans un espace vectoriel unifié.

Selon Google, Gemini Embedding 2 prend en charge :

• Texte avec jusqu'à 8192 tokens d'entrée
• Images, avec jusqu'à 6 images par requête au format PNG et JPEG
• Vidéo jusqu'à 120 secondes en mp4 et mov
• Audio sans nécessiter de transcription préalable
• Documents PDF jusqu'à 6 pages

Il prend également en charge les entrées multimodales entrelacées, comme une image plus du texte dans une seule requête. Cela permet au modèle de capturer des relations plus riches entre différents types de données.

Une autre fonctionnalité importante est la dimensionnalité de sortie flexible grâce à Matryoshka Representation Learning. La taille par défaut est de 3072 dimensions, mais elle peut être réduite à des tailles plus petites comme 1536 ou 768. Cela aide les développeurs à équilibrer qualité, stockage et vitesse de récupération en fonction de l'application.

Construire un système de correspondance d'images utilisant Gemini Embedding 2

Le projet utilise trois dossiers à l'intérieur d'un répertoire de données :

• dataset/nitika/vasu/janvi/

Chaque dossier contient plusieurs images d'une personne. L'objectif est simple :

• Lire toutes les images du dataset
• Générer un embedding pour chaque image en utilisant Gemini Embedding 2
• Stocker ces embeddings en mémoire et les mettre en cache localement
• Prendre une image de requête
• Générer son embedding
• Comparer cet embedding avec tous les embeddings d'images stockés en utilisant la similarité cosinus
• Retourner les images les plus correspondantes et prédire le nom de la personne

C'est un excellent exemple de la façon dont Gemini Embedding 2 peut être utilisé pour la récupération basée sur les images et la classification légère.

La meilleure partie de ce projet est qu'il ne nécessite pas de pipeline d'entraînement en apprentissage profond complet. Il n'y a pas d'entraînement de CNN personnalisé, pas de réglage fin et pas de flux de travail lourd en annotation. Au lieu de cela, le système s'appuie sur le modèle d'embedding en tant qu'extracteur de caractéristiques sémantiques.

Cela rend le développement beaucoup plus rapide.

Étant donné que Gemini Embedding 2 est nativement multimodal, la même conception de projet peut être étendue au-delà des images. Par exemple :

• Correspondre un clip audio parlé à un profil de personne
• Rechercher un PDF pertinent à partir d'une image
• Récupérer un segment vidéo à partir d'une requête textuelle
• Comparer des descriptions d'images et de textes mélangées dans un seul espace d'embedding

Dans ce sens, le projet actuel est un point d'entrée simple dans une architecture de récupération multimodale beaucoup plus large.

Utilisation de l'API Gemini Embedding 2

Google fournit le modèle Gemini Embedding 2 via l'API Gemini et Vertex AI. L'appel d'embedding se fait par la méthode embed_content.

Un exemple multimodal de Google ressemble à ceci :

from google import genai
from google.genai import types

client = genai.Client()

with open("example.png", "rb") as f:
    image_bytes = f.read()

with open("sample.mp3", "rb") as f:
    audio_bytes = f.read()

result = client.models.embed_content(
    model="gemini-embedding-2-preview",
    "What is the meaning of life?",
    types.Part.from_bytes(
        data=image_bytes,
        mime_type="image/png",
        types.Part.from_bytes(
            data=audio_bytes,
            mime_type="audio/mpeg",
        )
    )
)

print(result.embeddings)

Pour mon projet, j'avais seulement besoin de la partie image de ce flux de travail. Au lieu d'envoyer du texte, une image et de l'audio ensemble, j'ai utilisé une seule image par requête et généré son embedding.

Mise en œuvre du projet

Le projet commence par charger la clé API de Gemini à partir d'un fichier .env et créer un client :

from dotenv import load_dotenv
from google import genai

GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
client = genai.Client(api_key=GEMINI_API_KEY)

Ensuite, j'ai défini des fonctions d'assistance pour la validation des images, la détection du type MIME, la normalisation, la similarité cosinus et l'affichage des images.

La fonction principale d'embedding lit les octets de l'image et les envoie à Gemini Embedding 2 :

def embed_image(image_path):
    image_path = Path(image_path)
    mime_type = guess_mime_type(image_path)
    with open(image_path, "rb") as f:
        image_bytes = f.read()
    result = client.models.embed_content(
        model="gemini-embedding-2-preview",
        types.Part.from_bytes(
            data=image_bytes,
            mime_type=mime_type,
            config=types.EmbedContentConfig(
                output_dimensionality=3072
            )
        )
    )
    emb = np.array(result.embeddings[0].values, dtype=np.float32)
    return normalize(emb)

Cette fonction est le cœur de l'ensemble du pipeline. Elle transforme chaque image en une représentation vectorielle de 3072 dimensions.

Construction de la base de données d'embeddings du dataset

L'étape suivante consiste à parcourir le dossier du dataset, lire toutes les images pour chaque personne et les intégrer une par une.

Chaque image intégrée est stockée sous forme de dictionnaire contenant :

• le label de la personne
• le vecteur d'embedding

Pour éviter de recalculer les embeddings à chaque fois, je les ai mis en cache dans un fichier pickle local :

def build_embeddings_db(dataset, cache_file="image_embeddings_cache.pkl", force_rebuild=False):
    cache_path = Path(cache_file)
    if cache_path.exists() and not force_rebuild:
        with open(cache_path, "rb") as f:
            embeddings_db = pickle.load(f)
        return embeddings_db

    embeddings_db = []
    for item in dataset:
        emb = embed_image(item["path"])
        embeddings_db.append({
            "label": item["label"],
            "path": item["path"],
            "embedding": emb
        })
    
    with open(cache_path, "wb") as f:
        pickle.dump(embeddings_db, f)
    
    return embeddings_db

Cela rend le notebook beaucoup plus efficace car les embeddings ne sont générés qu'une seule fois, sauf si le dataset change.

Correspondre une image de requête

Une fois que les embeddings du dataset sont prêts, l'étape suivante consiste à tester le système avec une nouvelle image de requête.

L'image de requête est intégrée en utilisant la même fonction. Ensuite, son embedding est comparé à tous les embeddings stockés en utilisant la similarité cosinus.

def find_best_matches(query_image_path, top_k=5):
    query_emb = embed_image(query_image_path)
    results = []
    
    for item in embeddings_db:
        score = cosine_similarity(query_emb, item["embedding"])
        results.append({
            "label": item["label"],
            "path": item["path"],
            "score": score
        })
    
    results.sort(key=lambda x: x["score"], reverse=True)
    return results[:top_k]

Cette fonction retourne les meilleures images correspondantes du dataset.

Pour prédire le label final de la personne, j'ai utilisé le vote des top-k :

def predict_person(query_image_path, top_k=5):
    matches = find_best_matches(query_image_path, top_k=top_k)
    labels = [m["label"] for m in matches]
    predicted_label = Counter(labels).most_common(1)[0][0]
    return predicted_label, matches

C'est plus stable que de se fier à une seule image la plus proche.

Tester le projet

Dans le projet, j'ai testé des images de requête telles que :

query_image = "Nitika_Test_Image.jpeg"
predicted_person, matches = predict_person(query_image, top_k=2)
print("\nImage de requête :")
show_image(query_image, title="Image de requête")
print("Personne prédite :", predicted_person)
print("\nMeilleures correspondances :")
for i, match in enumerate(matches, 1):
    print(f"{i}. {match['label']} | score={match['score']:.4f} | path={match['path']}")
    show_image(match["path"], title=f"Rang {i} | {match['label']} | score={match['score']:.4f}")

Cela permet de visualiser les résultats et de vérifier la précision du système de correspondance d'images.