Multimodal Search through Image and Text

Use OpenAI's Clip neural network to embed images and text for vector similarity search

This example is aimed to guide you through using OpenAI's Clip model to embed a dataset, store it into LanceDB, and search for relevant texts/images. The "multimodal_clip" example also uses the same model.

First, let's install and import some dependencies. (These are all in the requirements.txt file in this example folder and the root folder.)

We will be using this HuggingFace dataset here, which contain pictures of various types of dogs and other animals.

Load the dataset

Unfortunately, this dataset only labels the images with numbers, so we can create an enum to map the numbers to the actual class names. The names are provided on the HuggingFace dataset link.

Enum class

We can now load the pretrained model ViT-B/32 from Clip, using either cuda or cpu depending on which torch version you are using.

Creating the image embedding function here. We want the embeddings to be a standard list, so we can convert the Tensor array to Numpy array to List.

Time to connect to the LanceDB table! you can create a PyArrow schema to initialize an empty table, like this:

	schema = pa.schema(
  [
      pa.field("vector", pa.list_(pa.float32(), 512)),
      pa.field("id", pa.int32()),
      pa.field("label", pa.int32()),
  ])
tbl = db.create_table("animal_images", schema=schema)

After creating the table, we want to prepare all the data to add to the table. We can first append all the data as a dictionary to an array.

Then, we can split it up into batches of 50, to then embed the image and add to the table. This will take around 10 minutes, because embedding 3.7k images can take a while.

The table looks good! We can now start testing the image search. First, load the validation split of the dataset.

This is how we can search the table: by embedding the image we want, calling the search function, and retuning a Pandas DataFrame.

The purpose of storing the ID is to call the original dataset back and displaying the image. Another way could be to store the image as bytes in the database, but this would make the database much larger.

And getting the id of the first result:

Nice! It correctly identified another italian greyhound! Now, to everything into a function:

Great, we have finished the image search! To begin the multimodal text search, we can create a similar embedding function, but using Clip's encode_text function instead of encode_image.

We perform the same search here by calling our just-made text embedder.

Combining everything into a function once again:

Great! It's working pretty well, except for the 4th image, which probably included the background into the embedding. But for now, it's good enough. You can always play around with this notebook and try to improve the results, but I hope you learned how to use multimodal search with OpenAI's Clip model!