Integrate image searchers

Image search allows searching for similar images in a database of images. It works by embedding the search query into a high-dimensional vector representing the semantic meaning of the query, followed by similarity search in a predefined, custom index using ScaNN (Scalable Nearest Neighbors).

As opposed to image classification, expanding the number of items that can be recognized doesn't require re-training the entire model. New items can be added simply re-building the index. This also enables working with larger (100k+ items) databases of images.

Use the Task Library ImageSearcher API to deploy your custom image searcher into your mobile apps.

Key features of the ImageSearcher API

  • Takes a single image as input, performs embedding extraction and nearest-neighbor search in the index.

  • Input image processing, including rotation, resizing, and color space conversion.

  • Region of interest of the input image.

Prerequisites

Before using the ImageSearcher API, an index needs to be built based on the custom corpus of images to search into. This can be achieved using Model Maker Searcher API by following and adapting the tutorial.

For this you will need:

After this step, you should have a standalone TFLite searcher model (e.g. mobilenet_v3_searcher.tflite), which is the original image embedder model with the index attached into the TFLite Model Metadata.

Run inference in Java

Step 1: Import Gradle dependency and other settings

Copy the .tflite searcher model file to the assets directory of the Android module where the model will be run. Specify that the file should not be compressed, and add the TensorFlow Lite library to the module’s build.gradle file:

android {
    // Other settings

    // Specify tflite index file should not be compressed for the app apk
    aaptOptions {
        noCompress "tflite"
    }

}

dependencies {
    // Other dependencies

    // Import the Task Vision Library dependency (NNAPI is included)
    implementation 'org.tensorflow:tensorflow-lite-task-vision:0.4.4'
    // Import the GPU delegate plugin Library for GPU inference
    implementation 'org.tensorflow:tensorflow-lite-gpu-delegate-plugin:0.4.4'
}

Step 2: Using the model

// Initialization
ImageSearcherOptions options =
    ImageSearcherOptions.builder()
        .setBaseOptions(BaseOptions.builder().useGpu().build())
        .setSearcherOptions(
            SearcherOptions.builder().setL2Normalize(true).build())
        .build();
ImageSearcher imageSearcher =
    ImageSearcher.createFromFileAndOptions(context, modelFile, options);

// Run inference
List<NearestNeighbor> results = imageSearcher.search(image);

See the source code and javadoc for more options to configure the ImageSearcher.

Run inference in C++

// Initialization
ImageSearcherOptions options;
options.mutable_base_options()->mutable_model_file()->set_file_name(model_path);
options.mutable_embedding_options()->set_l2_normalize(true);
std::unique_ptr<ImageSearcher> image_searcher = ImageSearcher::CreateFromOptions(options).value();

// Create input frame_buffer from your inputs, `image_data` and `image_dimension`.
// See more information here: tensorflow_lite_support/cc/task/vision/utils/frame_buffer_common_utils.h
std::unique_ptr<FrameBuffer> frame_buffer = CreateFromRgbRawBuffer(
      image_data, image_dimension);

// Run inference
const SearchResult result = image_searcher->Search(*frame_buffer).value();

See the source code for more options to configure ImageSearcher.

Run inference in Python

Step 1: Install TensorFlow Lite Support Pypi package.

You can install the TensorFlow Lite Support Pypi package using the following command:

pip install tflite-support

Step 2: Using the model

from tflite_support.task import vision

# Initialization
image_searcher = vision.ImageSearcher.create_from_file(model_path)

# Run inference
image = vision.TensorImage.create_from_file(image_file)
result = image_searcher.search(image)

See the source code for more options to configure ImageSearcher.

Example results

Results:
 Rank#0:
  metadata: burger
  distance: 0.13452
 Rank#1:
  metadata: car
  distance: 1.81935
 Rank#2:
  metadata: bird
  distance: 1.96617
 Rank#3:
  metadata: dog
  distance: 2.05610
 Rank#4:
  metadata: cat
  distance: 2.06347

Try out the simple CLI demo tool for ImageSearcher with your own model and test data.