Image Similarity with Hugging Face Datasets and Transformers

In this post, you'll learn to build an image similarity system with 🤗 Transformers. Finding out the similarity between a query image and potential candidates is an important use case for information retrieval systems, such as reverse image search, for example. All the system is trying to answer is that, given a query image and a set of candidate images, which images are the most similar to the query image.

We'll leverage the 🤗 datasets library as it seamlessly supports parallel processing which will come in handy when building this system.

Although the post uses a ViT-based model (nateraw/vit-base-beans) and a particular dataset (Beans), it can be extended to use other models supporting vision modality and other image datasets. Some notable models you could try:

• Swin Transformer
• ConvNeXT
• RegNet

Also, the approach presented in the post can potentially be extended to other modalities as well.

To study the fully working image-similarity system, you can refer to the Colab Notebook linked at the beginning.

How do we define similarity?

To build this system, we first need to define how we want to compute the similarity between two images. One widely popular practice is to compute dense representations (embeddings) of the given images and then use the cosine similarity metric to determine how similar the two images are.

For this post, we'll use “embeddings” to represent images in vector space. This gives us a nice way to meaningfully compress the high-dimensional pixel space of images (224 x 224 x 3, for example) to something much lower dimensional (768, for example). The primary advantage of doing this is the reduced computation time in the subsequent steps.

Computing embeddings

To compute the embeddings from the images, we'll use a vision model that has some understanding of how to represent the input images in the vector space. This type of model is also commonly referred to as image encoder.

Some questions that might arise here:

Q2: Why this checkpoint in particular?

Leads to: