Multimodal Similarity Search on the WikiArt Dataset

In this notebook, we build a similarity search tool capable of searching a dataset of text and images. For the dataset, we use WikiArt, which contains over 100k Western art paintings along with their titles, artists, and other metadata. For the embedding model, we use the well-known CLIP model, which can embed text and images into the same latent space. We will demonstrate how to search for a painting matching the query with the painting titles (text) and the thumbnails (images) separately, then using a reranker to combine results.

First, our imports:

1. Embed images and text into the same cross-modal space

(a) Load the CLIP multimodal embedding models

First, we load our image and text encoders, which are pre-trained neural networks that convert the image and text into 512-dimensional embeddings. These encoders are designed so that the image and text embeddings will be close to each other in space if they are semantically similar.

We use the CLIP model, which is freely available and can be easily loaded with the HuggingFace suite of libraries. For more details, refer to this notebook and this documentation. For the technically curious, the image encoder relies on a vision transformer (ViT), while the text encoder uses the same transformer decoder-only architecture as GPT-2, a precursor to ChatGPT.

(b) Open the WikiArt dataset

We will use a version of the WikiArt dataset available here, which is easily accessible through HuggingFace libraries. This dataset includes 103,250 rows, each representing a notable painting. For each painting, it provides a small image along with its title, artist, date, genre, and style. The dataset is widely known for its use in computer vision research.

The CLIP image encoder expects input images with a width of 224, while the dataset images have a width of 256. To address this, we will use Torch Vision to resize each image. We open the dataset in streaming mode and apply the map operation, which will only take effect when the data is iterated over, since the dataset is streaming. In other words, the map operation is lazy.

(I typically download the dataset in advance using git clone git@hf.co/datasets/Artificio/WikiArt to keep track of all my data.)

Let's examine a few samples to understand what's in the data:

Calculating embeddings for the image and title using the models we've loaded is a simple task:

This produces embeddings of size 512. The image encoder requires a list of PIL Images as input, while the text encoder takes a list of strings.

Since we can't reset the streaming iterator to the start, we'll create it again.

(c) Create Milvus vector database

So far, we've loaded our embedding models, the dataset we're embedding, and demonstrated how to embed images and text. As we calculate the embeddings for each painting, we will store them in a Milvus vector database. In this section, we will set up the database.

First, we create the database file on our local machine, define a schema for each row, and add a collection, similar to a table in a relational database. By opening a connection to a .db file, we're using Milvus Lite, a version suited for single machines and notebooks. It's worth noting that the API is nearly identical for Milvus Standalone, Milvus Distributed, and Zilliz CLoud.

We will store both the text embedding for each painting's caption and the corresponding image embedding in the same row, using a Milvus feature called multivectors. For convenience, we will store the image itself in each row, though we could reduce the database size by providing a filepath to each painting's JPEG, as these files are stored separately.

The database is now ready to accept new entries.

(d) Encode images and text and insert into database

The final step before we can search our painting dataset is to embed all the images and captions and build the database. We accomplish this using the following loop over wikiart_pt, which calculates the embeddings, makes minor adjustments to the metadata, and inserts batches into the database.

(Note that this inference loop is quite suboptimal. Much of the loop's time is spent compressing images to JPEG while the GPU remains idle. On my machine with an RTX 4090, this process takes about 21 minutes, but it would take only about 3 minutes without image compression. In the future, I plan to create a new dataset where resizing and image compression are already completed and uploaded to HuggingFace. This would simplify the tutorial significantly.)

The embeddings are normalized so that the dot product (also known as the inner product) between them equals the cosine distance.

We wrap the HuggingFace IterableDataset in a PyTorch DataLoader to leverage multiple CPUs for preprocessing the images. However, this approach has a downside: PIL Images get converted into tensors, requiring us to convert them back to PIL format. Subclassing DataLoader might offer a solution to this issue.

If everything goes as planned, our database will contain 103,250 rows.

We have successfully constructed our database and can now perform some interesting searches.

2. Search vector database

Let's create helper functions to search our database. We use a feature of Milvus called hybrid search to search separate indices and combine the results into a single set. We combine the results using a re-ranker, specifically the reciprocal rank fusion method.

The function searchByText() takes a text string as the input query and performs several types of searches. You can conduct a hybrid search across both text and image embeddings or search using just one of them. Additionally, you can pass in a filtering expression to filter on the metadata.

(a) Combining image and text vector search vs searching each separately

What benefits come from embedding both a text and an image of each painting? Why not just embed the image and perform a search on it? Let's compare the results of a search query using both indices, only the text index, and only the image index.

Here are the results for the query 'the starry night' searching both images and text:

Here are the same results but searching just the titles:

You can see that all the titles of the resulting images resemble the query. Interestingly, our CLIP model, which is multilingual, can embed the Swedish phrase "den ljusa natten," translating to "the bright night."

Let's compare now to searching just the images.

The results include paintings featuring a night sky with stars or the aurora borealis, which cannot be fully inferred from the title alone. While these results are sensible, they do not include Van Gogh's impressionist painting The Starry Night. In this instance, the title aids the search by clarifying the abstract representation of stars in the painting. This example demonstrates how combining image and text in hybrid search can yield more relevant results than using either method alone.

(b) Filtering search on metadata

Let's perform some searches now to demonstrate filtering by metadata. Suppose we want to find paintings from Picasso's Blue Period. We will search for the term "blue" and filter by the artist field:

This method seems to have successfully found paintings from the Blue Period.

To find paintings from Picasso's Cubist Period, you can search for the term "cubist" and filter by the artist field:

Once again, we can search for a style or theme of painting by a given artist in our dataset.

For another example, we'll search for Picasso's paintings that examine the theme of love.