Reverse Image Search featuring Impressionist Paintings

It's easy to get your own reverse-image search up and running locally with a vector database. We use Milvus Lite and PyTorch to build a local reverse image search using the Impressionist-Classifier Dataset found on Kaggle.

We start by importing the necessary libraries:

Next, we download the dataset using gdown to download it from a publicly hosted Google Drive and zipfile to unzip the images.

Once we've downloaded all the images, we need to set up some parameters that we'll use later on to work with our vector database.

Let's start our local Vector Database instance with Milvus Lite

With our server up, we're ready to get started. We start by defining the vector database schema and establishing a collection. Each entry in our collection features three fields. First, an id for regular querying, next a filepath to identify where the image is stored locally, and lastly, the embedding that we use for similarity search.

Once we have a connection to our vector database and an established collection, we create an index to search on. For this example, we use an IVF Flat index measured with the L2 Norm and 128 cluster units (nlist).

Click here to learn more about vector indexes.

Now let's get the embeddings. We use the ResNet50 model from PyTorch to get the embeddings. Normally, the last layer of the ResNet50 model outputs classifications for a dataset. In our case, we need the embeddings, not the classifications, so we remove the last layer.

At this point we have a local Milvus instance running, the data for this example project, and the model we need to get some embeddings. Uploading the data into the vector database as a collection and getting it indexed is the next step. First, we preprocess the data to fit the data that the ResNet50 Model was trained on (seen on the bottom of the page). Then we batch the data and "upload" it to our vector database.

the second block of this step takes a while to run (about 12 minutes on a 16GB RAM M1 Mac), we are running almost 5000 images through ResNet50, now would be a good time to grab a snack or something to drink :)

With our vector database populated, we are ready to perform a reverse image search. The second block in this section shows how you pick an image or set of images to reverse image search. For this example, we have two search patterns. The first one (commented out) shows how to search the provided test paintings. The second one shows how to reverse image search for a specific painting provided in the training set. In most cases, it's bad practice to use training data while doing any sort of validation, but this particular case confirms that we are returning the same image when performing a reverse image search an on image that exists in our vector database

We can plot the images for a visual.

We can also see which images were classified as the most similar.

Finally, we want to clean up our Milvus instance.