Image Search

Learn how to perform similarity search over images, generating image embeddings and using TileDB-Vector-Search to add them to vector indexes and query them.

How to run this tutorial

We recommend running this tutorial, as well as the other various tutorials in the Tutorials section, inside TileDB Cloud. This will allow you to quickly experiment avoiding all the installation, deployment, and configuration hassles. Sign up for the free tier, spin up a TileDB Cloud notebook with a Python kernel, and follow the tutorial instructions. If you wish to learn how to run tutorials locally on your machine, read the Tutorials: Running Locally tutorial.

This tutorial is an image search demo, using TileDB-Vector-Search and a dataset of flowers.

Setup

If you are running this tutorial on a local machine, you may have to perform the following installations:

conda install -y -c tiledb -c conda-forge -c anaconda tiledb-vector-search scikit-learn tensorflow-datasets matplotlib efficientnet
# NOTE: `efficientnet` conda package is available on linux only, so
# you can install it from pip on all non-Linux machines

# Or, using pip
#pip install tiledb tiledb-vector-search-scikit-learn tensorflow-datasets efficientnet matplotlib
pip install

Import the necessary packages, set the URIs you will use throughout this tutorial, and delete any past data:

import os
import shutil
import warnings

import numpy as np

warnings.filterwarnings("ignore")
import random

import PIL
import tensorflow as tf
import tensorflow_datasets as tfds
import tiledb
import tiledb.vector_search as vs
from efficientnet.preprocessing import center_crop_and_resize
from tensorflow.keras.applications.resnet_v2 import preprocess_input
from tensorflow_datasets.core import dataset_utils

# URIs you will use in this tutorial
dataset = "tf_flowers"
image_array_uri = "tf_flowers_array"
index_uri = "tf_flowers_index"
features_uri = "features.f32bin"

# Clean up past data
if os.path.exists(image_array_uri):
    shutil.rmtree(image_array_uri)
if os.path.exists(index_uri):
    shutil.rmtree(index_uri)
if os.path.exists(features_uri):
    os.remove(features_uri)

Store raw images

In addition to the vector embeddings you will store using TileDB-Vector-Search, you can use the TileDB array engine to store all the raw images in a 3D TileDB array. This is yet another cool feature of TileDB, which allows you to unify all your data (both “unstructured” and “structured”) under a common data model (i.e., arrays) and database engine.

Create the 3D image array as follows:

# Will be used to define the array domain
crop_image_size = 224

# Define the image array schema, which is a 3D array that will store all images
image_array_schema = tiledb.ArraySchema(
    domain=tiledb.Domain(
        [
            tiledb.Dim(
                name="image_id",
                dtype="uint64",
                domain=(0, np.iinfo(np.uint64).max - 10000),
                tile=10,
            ),
            tiledb.Dim(
                name="d1",
                dtype="uint64",
                domain=(0, crop_image_size - 1),
                tile=crop_image_size,
            ),
            tiledb.Dim(
                name="d2",
                dtype="uint64",
                domain=(0, crop_image_size - 1),
                tile=crop_image_size,
            ),
            tiledb.Dim(name="d3", dtype="uint64", domain=(0, 2), tile=2),
        ]
    ),
    attrs=[
        tiledb.Attr(name="value", dtype=np.uint8),
    ],
    sparse=False,
)

# Physically create the array
tiledb.Array.create(image_array_uri, image_array_schema)

Next, load the images from their source location.

# Load the image dataset.
# NOTE: This could take several minutes depending on your machine configuration.
ds, ds_info = tfds.load(dataset, split="train", with_info=True)

Create a utility function that loads and preprocesses the images and saves them in a TileDB array:

def save_images(ds_info, ds, num_samples, image_key, crop_image_size=-1):
    samples = list(dataset_utils.as_numpy(ds.take(num_samples)))
    images_data = np.array([])
    for i, sample in enumerate(samples):
        if not isinstance(sample, dict):
            raise ValueError(
                "tfds.show_examples requires examples as `dict`, with the same "
                "structure as `ds_info.features`. It is currently not compatible "
                "with `as_supervised=True`. Received: {}".format(type(sample))
            )

        # Preprocess the image
        image = sample[image_key]
        if len(image.shape) != 3:
            raise ValueError(
                "Image dimension should be 3. tfds.show_examples does not support "
                "batched examples or video."
            )
        _, _, c = image.shape
        if c == 1:
            image = image.reshape(image.shape[:2])
        if crop_image_size != -1:
            image = center_crop_and_resize(image, crop_image_size).astype(np.uint8)
        if images_data.any():
            images_data = np.concatenate((images_data, image[None, :]), axis=0)
        else:
            images_data = image[None, :]
    with tiledb.open(image_array_uri, mode="w") as A:
        A[0:image_samples] = {"value": images_data}

Now, save the images in the TileDB array:

image_samples = 3670
save_images(ds_info, ds, image_samples, "image", crop_image_size)

Create vector index

Create a utility function that computes the image embeddings using ResNet50V2:

def calculate_resnet(x: np.ndarray) -> np.ndarray:
    model = tf.keras.applications.ResNet50V2(include_top=False)
    maps = model.predict(preprocess_input(x))
    if np.prod(maps.shape) == maps.shape[-1] * len(x):
        return np.squeeze(maps)
    else:
        return maps.mean(axis=1).mean(axis=1)

Next, generate embeddings for all images and store them in a binary file:

with tiledb.open(image_array_uri, mode="r") as A:
    data = A[0:image_samples]["value"]
    embeddings = calculate_resnet(data)
    with open("features.f32bin", "wb") as f:
        np.array(embeddings.shape, dtype="uint32").tofile(f)
        np.array(embeddings).astype("float32").tofile(f)

Ingest the embeddings into an IVF_FLAT array:

index_ivf_flat = vs.ingest(
    index_type="IVF_FLAT",
    index_uri=index_uri,
    source_uri="features.f32bin",
)

Similarity search

Prepare the index for reading:

index_flat = vs.IVFFlatIndex(index_uri)

Choose a randomly sampled image from the image dataset:

with tiledb.open(image_array_uri, mode="r") as A:
    query_image = A[random.randint(0, image_samples - 1)]["value"]
query_embeding = calculate_resnet(query_image[None, :])
display(PIL.Image.fromarray(query_image))

1/1 [==============================] - 1s 616ms/step

Perform image similarity search using the selected image:

# Retrieve the 5 most similar images to the randomly selected one.
result_d, result_i = index_flat.query(query_embeding, k=5)
with tiledb.open(image_array_uri, mode="r") as A:
    for result in result_i[0]:
        display(PIL.Image.fromarray(A[result]["value"]))

Clean up

Delete all the generated data.

# Clean up
if os.path.exists(image_array_uri):
    shutil.rmtree(image_array_uri)
if os.path.exists(index_uri):
    shutil.rmtree(index_uri)
if os.path.exists(features_uri):
    os.remove(features_uri)