Similitud de imágenes con conjuntos de datos y transformadores de Hugging Face

Imagenes similares con conjuntos de datos y Transformers de Hugging Face.

En esta publicación, aprenderás a construir un sistema de similitud de imágenes con 🤗 Transformers. Descubrir la similitud entre una imagen de consulta y posibles candidatas es un caso de uso importante para los sistemas de recuperación de información, como la búsqueda inversa de imágenes, por ejemplo. Todo lo que el sistema intenta responder es qué imágenes son las más similares a la imagen de consulta, dada una imagen de consulta y un conjunto de imágenes candidatas.

Aprovecharemos la biblioteca 🤗 datasets ya que admite el procesamiento paralelo, lo cual será útil al construir este sistema.

Aunque la publicación utiliza un modelo basado en ViT (nateraw/vit-base-beans) y un conjunto de datos en particular (Beans), se puede ampliar para usar otros modelos que admitan la modalidad de visión y otros conjuntos de datos de imágenes. Algunos modelos destacados que podrías probar:

  • Swin Transformer
  • ConvNeXT
  • RegNet

Además, el enfoque presentado en la publicación se puede extender potencialmente a otras modalidades también.

Para estudiar el sistema de similitud de imágenes completamente funcional, puedes consultar el Cuaderno de Colab vinculado al principio.

¿Cómo definimos la similitud?

Para construir este sistema, primero debemos definir cómo queremos calcular la similitud entre dos imágenes. Una práctica ampliamente popular es calcular representaciones densas (incrustaciones) de las imágenes dadas y luego utilizar la métrica de similitud del coseno para determinar qué tan similares son las dos imágenes.

Para esta publicación, usaremos “incrustaciones” para representar imágenes en un espacio vectorial. Esto nos brinda una forma interesante de comprimir de manera significativa el espacio de píxeles de alta dimensión de las imágenes (por ejemplo, 224 x 224 x 3) a algo mucho más dimensionalmente bajo (768, por ejemplo). La ventaja principal de hacer esto es el tiempo de cálculo reducido en los pasos posteriores.

Cálculo de las incrustaciones

Para calcular las incrustaciones de las imágenes, utilizaremos un modelo de visión que tiene cierta comprensión de cómo representar las imágenes de entrada en el espacio vectorial. Este tipo de modelo también se conoce comúnmente como codificador de imágenes.

Para cargar el modelo, aprovechamos la clase AutoModel. Proporciona una interfaz para cargar cualquier punto de control de modelo compatible desde el Hugging Face Hub. Junto con el modelo, también cargamos el procesador asociado con el modelo para la preprocesamiento de datos.

from transformers import AutoFeatureExtractor, AutoModel


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

En este caso, el punto de control se obtuvo mediante el ajuste fino de un modelo basado en Vision Transformer en el conjunto de datos beans.

Algunas preguntas que podrían surgir aquí:

P1 : ¿Por qué no usamos AutoModelForImageClassification?

Esto se debe a que queremos obtener representaciones densas de las imágenes y no categorías discretas, que es lo que AutoModelForImageClassification proporcionaría.

P2 : ¿Por qué este punto de control en particular?

Como se mencionó anteriormente, estamos utilizando un conjunto de datos específico para construir el sistema. Entonces, en lugar de usar un modelo generalista (como los entrenados en el conjunto de datos ImageNet-1k, por ejemplo), es mejor usar un modelo que se haya ajustado finamente al conjunto de datos utilizado. De esta manera, el modelo subyacente comprende mejor las imágenes de entrada.

Ten en cuenta que también puedes usar un punto de control que se haya obtenido mediante el preentrenamiento auto-supervisado. El punto de control no necesariamente tiene que provenir del aprendizaje supervisado. De hecho, si se pre-entrenan correctamente, los modelos auto-supervisados pueden ofrecer un rendimiento de recuperación impresionante.

Ahora que tenemos un modelo para calcular las incrustaciones, necesitamos algunas imágenes candidatas para hacer consultas.

Cargando un conjunto de datos para imágenes candidatas

En algún momento, construiremos tablas hash que mapean las imágenes candidatas a hashes. Durante el tiempo de consulta, utilizaremos estas tablas hash. Hablaremos más sobre las tablas hash en la sección respectiva, pero por ahora, para tener un conjunto de imágenes candidatas, usaremos la división train del conjunto de datos beans.

from datasets import load_dataset


dataset = load_dataset("beans")

Así es como se ve una muestra única de la división de entrenamiento:

El conjunto de datos tiene tres características:

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

Para demostrar el sistema de similitud de imágenes, utilizaremos 100 muestras del conjunto de datos de imágenes candidatas para acortar el tiempo de ejecución general.

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

El proceso de búsqueda de imágenes similares

A continuación, se puede encontrar una descripción general pictórica del proceso subyacente para obtener imágenes similares.

Desglosando un poco la figura anterior, tenemos:

  1. Extraer los embeddings de las imágenes candidatas (candidate_subset), almacenándolos en una matriz.
  2. Tomar una imagen de consulta y extraer sus embeddings.
  3. Iterar sobre la matriz de embeddings (calculada en el paso 1) y calcular el puntaje de similitud entre el embedding de consulta y los embeddings de candidatos actuales. Normalmente, mantenemos una correspondencia similar a un diccionario que mantiene una correspondencia entre algún identificador de la imagen candidata y los puntajes de similitud.
  4. Ordenar la estructura de mapeo con respecto a los puntajes de similitud y devolver los identificadores subyacentes. Utilizamos estos identificadores para obtener las muestras candidatas.

Podemos escribir una utilidad simple y aplicarle map() a nuestro conjunto de datos de imágenes candidatas para calcular los embeddings de manera eficiente.

import torch 

def extract_embeddings(model: torch.nn.Module):
    """Utilidad para calcular los embeddings."""
    device = model.device

    def pp(batch):
        images = batch["image"]
        # `transformation_chain` es una composición de transformaciones de preprocesamiento
        # que aplicamos a las imágenes de entrada para prepararlas
        # para el modelo. Para más detalles, consulte el Cuaderno de Colab adjunto.
        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

Y podemos aplicar extract_embeddings() de la siguiente manera:

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)

A continuación, para mayor comodidad, creamos una lista que contiene los identificadores de las imágenes candidatas.

candidate_ids = []

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

    # Crear un identificador único.
    entry = str(id) + "_" + str(label)

    candidate_ids.append(entry)

Utilizaremos la matriz de los embeddings de todas las imágenes candidatas para calcular los puntajes de similitud con una imagen de consulta. Ya hemos calculado los embeddings de las imágenes candidatas. En la siguiente celda, simplemente los juntamos en una matriz.

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

Utilizaremos la similitud del coseno para calcular el puntaje de similitud entre dos vectores de embedding. Luego lo utilizaremos para obtener muestras candidatas similares dada una muestra de consulta.

def compute_scores(emb_one, emb_two):
    """Calcula la similitud del coseno entre dos vectores."""
    scores = torch.nn.functional.cosine_similarity(emb_one, emb_two)
    return scores.numpy().tolist()


def fetch_similar(image, top_k=5):
    """Obtiene las `top_k` imágenes similares con `image` como consulta."""
    # Prepara la imagen de consulta de entrada para el cálculo del embedding.
    image_transformed = transformation_chain(image).unsqueeze(0)
    new_batch = {"pixel_values": image_transformed.to(device)}

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

    # Calcula los puntajes de similitud con todas las imágenes candidatas de una vez.
    # También creamos un mapeo entre los identificadores de imágenes candidatas
    # y sus puntajes de similitud con la imagen de consulta.
    sim_scores = compute_scores(all_candidate_embeddings, query_embeddings)
    similarity_mapping = dict(zip(candidate_ids, sim_scores))
 
    # Ordena el diccionario de mapeo y devuelve las `top_k` candidatas.
    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

Realizar una consulta

Dadas todas las utilidades, estamos preparados para realizar una búsqueda de similitud. Vamos a tener una imagen de consulta del conjunto de prueba (test) del conjunto de datos 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"Etiqueta de la consulta: {test_label}")
print(f"Etiquetas de los 5 candidatos principales: {sim_labels}")

Conduce a:

Etiqueta de la consulta: 0
Etiquetas de los 5 candidatos principales: [0, 0, 0, 0, 0]

Parece que nuestro sistema obtuvo el conjunto correcto de imágenes similares. Cuando se visualizan, obtendríamos:

Extensiones adicionales y conclusiones

Ahora tenemos un sistema de similitud de imágenes funcional. Pero en realidad, estarás lidiando con muchas más imágenes candidatas. Teniendo eso en cuenta, nuestro procedimiento actual tiene múltiples inconvenientes:

  • Si almacenamos los embeddings tal cual, los requisitos de memoria pueden aumentar rápidamente, especialmente cuando se trata de millones de imágenes candidatas. Los embeddings son de 768-d en nuestro caso, lo que aún puede ser relativamente alto en el régimen a gran escala.
  • Tener embeddings de alta dimensionalidad tiene un efecto directo en los cálculos posteriores involucrados en la parte de recuperación.

Si de alguna manera podemos reducir la dimensionalidad de los embeddings sin alterar su significado, aún podemos mantener un buen equilibrio entre velocidad y calidad de recuperación. El Cuaderno de Colab adjunto a esta publicación implementa y demuestra utilidades para lograr esto con proyección aleatoria y hashing sensible a la localidad.

🤗 Datasets ofrece integraciones directas con FAISS, lo que simplifica aún más el proceso de construcción de sistemas de similitud. Supongamos que ya has extraído los embeddings de las imágenes candidatas (el conjunto de datos beans) y los has almacenado dentro de una característica llamada embeddings. Ahora puedes usar fácilmente el método add_faiss_index() del conjunto de datos para construir un índice denso:

dataset_with_embeddings.add_faiss_index(column="embeddings")

Una vez construido el índice, se puede utilizar dataset_with_embeddings para recuperar los ejemplos más cercanos dados los embeddings de consulta con get_nearest_examples():

puntuaciones, ejemplos_recuperados = dataset_with_embeddings.get_nearest_examples(
    "embeddings", embeddings_de_consulta, k=top_k
)

El método devuelve las puntuaciones y los ejemplos candidatos correspondientes. Para obtener más información, puedes consultar la documentación oficial y este cuaderno.

Finalmente, puedes probar el siguiente Espacio que construye una aplicación mini de similitud de imágenes:

En esta publicación, repasamos un inicio rápido para construir sistemas de similitud de imágenes. Si encontraste esta publicación interesante, recomendamos encarecidamente construir sobre los conceptos que discutimos aquí para que te sientas más cómodo con el funcionamiento interno.

¿Sigues buscando aprender más? Aquí hay algunos recursos adicionales que pueden ser útiles para ti:

  • Faiss: una biblioteca para búsqueda de similitud eficiente
  • ScaNN: búsqueda eficiente de similitud de vectores
  • Integración de buscadores de imágenes en aplicaciones móviles