Creación de un motor de búsqueda de similitud de imágenes con FAISS y CLIP | por Lihi Gur Arie, PhD | agosto de 2024

Un tutorial guiado que explica cómo buscar en su conjunto de datos de imágenes con consultas de texto o fotos, utilizando la incrustación CLIP y la indexación FAISS.

La imagen fue generada por el autor en la plataforma Flux-Pro

¿Alguna vez ha deseado encontrar una imagen entre su interminable conjunto de datos de imágenes, pero le resultó demasiado tedioso? En este tutorial, crearemos un motor de búsqueda de similitud de imágenes para encontrar imágenes fácilmente mediante una consulta de texto o una imagen de referencia. Para su comodidad, el código completo de este tutorial se proporciona al final del artículo como un Cuaderno de Colab.

Descripción general del pipeline

El significado semántico de una imagen se puede representar mediante un vector numérico denominado incrustación. La comparación de estos vectores de incrustación de baja dimensión, en lugar de las imágenes sin procesar, permite realizar búsquedas de similitud eficientes. Para cada imagen del conjunto de datos, crearemos un vector de incrustación y lo almacenaremos en un índice. Cuando se proporciona una consulta de texto o una imagen de referencia, se genera su incrustación y se compara con las incrustaciones indexadas para recuperar las imágenes más similares.

He aquí una breve descripción general:

  1. Incrustar: Las incrustaciones de las imágenes se extraen utilizando el modelo CLIP.
  2. Indexación:Las incrustaciones se almacenan como un índice FAISS.
  3. Recuperación:Con FAISS, la incrustación de la consulta se compara con las incrustaciones indexadas para recuperar las imágenes más similares.

Modelo CLIP

El modelo CLIP (Contrastive Language-Image Pre-training), desarrollado por OpenAI, es un modelo de visión y lenguaje multimodal que asigna imágenes y texto al mismo espacio latente. Dado que utilizaremos consultas de imágenes y texto para buscar imágenes, utilizaremos el modelo CLIP para integrar nuestros datos. Para obtener más información sobre CLIP, puede consultar mi artículo anterior aquí.

Índice FAISS

FAISS (Facebook AI Similarity Search) es una biblioteca de código abierto desarrollada por Meta. Está construida alrededor del objeto Index que almacena los vectores de incrustación de la base de datos. FAISS permite una búsqueda de similitud eficiente y la agrupación de vectores densos, y la usaremos para indexar nuestro conjunto de datos y recuperar las fotos que se parecen a la consulta.

Paso 1: Exploración del conjunto de datos

Para crear el conjunto de datos de imágenes para este tutorial, recopilé 52 imágenes de diversos temas de PexelsPara tener una idea, observemos 10 imágenes al azar:

Paso 2: Extraer incrustaciones CLIP del conjunto de datos de imágenes

Para extraer incrustaciones CLIP, primero cargaremos el modelo CLIP usando la biblioteca HuggingFace SentenceTransformer:

model = SentenceTransformer('clip-ViT-B-32')

A continuación, crearemos una función que itere a través de nuestro directorio de conjunto de datos con globabre cada imagen con PIL Image.openy genera un vector de incrustación para cada imagen con CLIP model.encodeDevuelve una lista de los vectores de incrustación y una lista de las rutas de nuestro conjunto de datos de imágenes:

def generate_clip_embeddings(images_path, model):

image_paths = glob(os.path.join(images_path, '**/*.jpg'), recursive=True)

embeddings = []
for img_path in image_paths:
image = Image.open(img_path)
embedding = model.encode(image)
embeddings.append(embedding)

return embeddings, image_paths

IMAGES_PATH = '/path/to/images/dataset'

embeddings, image_paths = generate_clip_embeddings(IMAGES_PATH, model)

Paso 3: Generar el índice FAISS

El siguiente paso es crear un índice FAISS a partir de la lista de vectores de incrustación. FAISS ofrece varias métricas de distancia para la búsqueda de similitud, incluidas la distancia L2 (euclidiana) y el producto interno (IP).

FAISS también ofrece varias opciones de indexación. Puede utilizar técnicas de aproximación o compresión para manejar conjuntos de datos grandes de manera eficiente y, al mismo tiempo, equilibrar la velocidad y la precisión de la búsqueda. En este tutorial, utilizaremos un índice “plano”, que realiza una búsqueda de fuerza bruta comparando el vector de consulta con cada uno de los vectores del conjunto de datos, lo que garantiza resultados exactos a costa de una mayor complejidad computacional.

def create_faiss_index(embeddings, image_paths, output_path):

dimension = len(embeddings[0])
index = faiss.IndexFlatIP(dimension)
index = faiss.IndexIDMap(index)

vectors = np.array(embeddings).astype(np.float32)

# Add vectors to the index with IDs
index.add_with_ids(vectors, np.array(range(len(embeddings))))

# Save the index
faiss.write_index(index, output_path)
print(f"Index created and saved to {output_path}")

# Save image paths
with open(output_path + '.paths', 'w') as f:
for img_path in image_paths:
f.write(img_path + '\n')

return index

OUTPUT_INDEX_PATH = "/content/vector.index"
index = create_faiss_index(embeddings, image_paths, OUTPUT_INDEX_PATH)

El faiss.IndexFlatIP inicializa un índice para la similitud del producto interno, envuelto en un faiss.IndexIDMap para asociar cada vector con un ID. A continuación, el index.add_with_ids agrega los vectores al índice con identificaciones secuenciales y el índice se guarda en el disco junto con las rutas de las imágenes.

El índice se puede utilizar inmediatamente o guardar en el disco para uso futuro. Para cargar el índice FAISS usaremos esta función:

def load_faiss_index(index_path):
index = faiss.read_index(index_path)
with open(index_path + '.paths', 'r') as f:
image_paths = [line.strip() for line in f]
print(f"Index loaded from {index_path}")
return index, image_paths

index, image_paths = load_faiss_index(OUTPUT_INDEX_PATH)

Paso 4: Recuperar imágenes mediante una consulta de texto o una imagen de referencia

Con nuestro índice FAISS creado, ahora podemos recuperar imágenes mediante consultas de texto o imágenes de referencia. Si la consulta es una ruta de imagen, la consulta se abre con PIL Image.openA continuación, se extrae el vector de incrustación de la consulta con CLIP model.encode.

def retrieve_similar_images(query, model, index, image_paths, top_k=3):

# query preprocess:
if query.endswith(('.png', '.jpg', '.jpeg', '.tiff', '.bmp', '.gif')):
query = Image.open(query)

query_features = model.encode(query)
query_features = query_features.astype(np.float32).reshape(1, -1)

distances, indices = index.search(query_features, top_k)

retrieved_images = [image_paths[int(idx)] for idx in indices[0]]

return query, retrieved_images

La recuperación se está llevando a cabo en el index.search método. Implementa una búsqueda de k vecinos más cercanos (kNN) para encontrar el k Vectores más similares al vector de consulta. Podemos ajustar el valor de k cambiando el top_k parámetro. La métrica de distancia utilizada en la búsqueda kNN en nuestra implementación es la similitud del coseno. La función devuelve la consulta y una lista de rutas de imágenes recuperadas.

Buscar con una consulta de texto:

Ahora estamos listos para examinar los resultados de la búsqueda. La función auxiliar visualize_results Muestra los resultados. Puedes encontrarlos en el cuaderno de Colab asociado. Exploremos las 3 imágenes más similares recuperadas para la consulta de texto “pelota”, por ejemplo:

query = 'ball'
query, retrieved_images = retrieve_similar_images(query, model, index, image_paths, top_k=3)
visualize_results(query, retrieved_images)
Imágenes recuperadas con la consulta: ‘una pelota’

Para la consulta ‘animal’ obtenemos:

Imágenes recuperadas con la consulta: ‘animal’

Buscar con una imagen de referencia:

query ='/content/drive/MyDrive/Colab Notebooks/my_medium_projects/Image_similarity_search/image_dataset/pexels-w-w-299285-889839.jpg'
query, retrieved_images = retrieve_similar_images(query, model, index, image_paths, top_k=3)
visualize_results(query, retrieved_images)
Consulta e imágenes recuperadas

Como podemos ver, obtenemos resultados bastante interesantes para un modelo preentrenado listo para usar. Cuando realizamos una búsqueda por una imagen de referencia de una pintura de ojos, además de encontrar la imagen original, encontramos una coincidencia de anteojos y otra de una pintura diferente. Esto demuestra diferentes aspectos del significado semántico de la imagen de consulta.

Puede probar otras consultas en el cuaderno Colab proporcionado para ver cómo funciona el modelo con diferentes entradas de texto e imágenes.

En este tutorial, creamos un motor de búsqueda de similitud de imágenes básico utilizando CLIP y FAISS. Las imágenes recuperadas compartían un significado semántico similar con la consulta, lo que indica la eficacia del enfoque. Aunque CLIP muestra buenos resultados para un modelo Zero Shot, puede exhibir un bajo rendimiento en datos fuera de distribución, tareas de grano fino y heredar el sesgo natural de los datos con los que se entrenó. Para superar estas limitaciones, puede probar otros modelos preentrenados similares a CLIP como en Clip abiertoo ajuste CLIP en su propio conjunto de datos personalizado.

Felicitaciones por haber llegado hasta aquí. Haz clic en 👍 para mostrar tu agradecimiento y aumentar la autoestima del algoritmo 🤓

¿Quieres saber más?

Cuaderno de Colab Enlace