Image search with 🤗 datasets

🤗 datasets is a library that makes it easy to access and share datasets. It also makes it easy to process data efficiently -- including working with data which doesn't fit into memory.

When datasets was first launched, it was associated mostly with text data. However, recently, datasets has added increased support for audio as well as images. In particular, there is now a datasets feature type for images. A previous blog post showed how datasets can be used with 🤗 transformers to train an image classification model. In this blog post, we'll see how we can combine datasets and a few other libraries to create an image search application.

First, we'll install datasets. Since we're going to be working with images, we'll also install pillow. We'll also need sentence_transformers and faiss. We'll introduce those in more detail below. We also install rich - we'll only briefly use it here, but it's a super handy package to have around -- I'd really recommend exploring it further!

!pip install datasets pillow rich faiss-gpu sentence_transformers 

To start, let's take a look at the image feature. We can use the wonderful rich library to poke around python objects (functions, classes etc.)

from rich import inspect
import datasets

inspect(datasets.Image, help=True)

We can see there a few different ways in which we can pass in our images. We'll come back to this in a little while.

A really nice feature of the datasets library (beyond the functionality for processing data, memory mapping etc.) is that you get some nice things 'for free'. One of these is the ability to add a faiss index to a dataset. faiss is a "library for efficient similarity search and clustering of dense vectors".

The datasets docs shows an example of using a faiss index for text retrieval. In this post we'll see if we can do the same for images.

The dataset: "Digitised Books - Images identified as Embellishments. c. 1510 - c. 1900"

Creating our dataset

Let's see what we get back.

dataset

dataset.features['label']