Combine AI Models for Image Search using ONNX and Milvus

Open Neural Network Exchange (ONNX) is an open format built to represent machine learning models. Since it was open-sourced in 2017, ONNX has developed into a standard for AI, providing building blocks for machine learning and deep learning models. ONNX defines a common file format to enable AI developers to use models with various frameworks, tools, runtimes, and compilers, and helps increase the speed of innovation in the artificial intelligence community.

Milvus is an open-source vector database that is highly flexible, reliable and blazing fast. It supports adding, deleting, updating and near real-time search of vector. Milvus has a comprehensive set of intuitive APIs, and support for multiple widely adopted index libraries (e.g. Faiss, NMSLIB, and Annoy), simplifying the index selection for a given scenario. Milvus is simple to use, and it has been used in hundreds of organizations and institutions worldwide, including image, audio and video search, recommendation, chatbot, new drug search, etc.

This article will introduce you how to use multiple models for image search based on ONNX and Milvus. It takes VGG16 and ResNet50 models as examples, uses ONNX to run different AI models to generate feature vectors, and finally performs feature vector retrieval in Milvus to return similar images.

Process Models with ONNX

The ONNX format can be easily exchanged between AI models. For example, the TensorFlow model can be converted to ONNX format and run in the Caffe environment. In this example, we convert the pre-trained ResNet50 model under the Keras framework to the ONNX format, and then call the VGG16 model in ONNX format to analyze different models.

from keras.applications.resnet50 import ResNet50
import tensorflow as tf

# load keras-resnet50 model and save as a floder
model_resnet50 = ResNet50(include_top=False, pooling='max', weights='imagenet')
tf.saved_model.save(model_resnet50, "keras_resnet50_model")

# convert resnet50 model to onnx
! python -m tf2onnx.convert --saved-model "keras_resnet50_model" --output "onnx_resnet50.onnx"

Note: When we used the interface keras2onnx.convert_keras(model, model.name) to convert the model, it will return the error AttributeError:'KerasTensor' object has no attribute'graph'. Then we can use Python’s Bash command to convert according to the solution on Stack Overflow.

Extract Feature Vectors using Models

After converting ResNet50 model into ONNX format, you can extract the feature vector of the picture directly through the inference. Note: Feature vectors need to be normalized after extraction.

# get the image vectors with onnx model
def get_onnx_vectors(onnx_model, img_path):
    img = image.load_img(img_path, target_size=(224, 224))
    x = image.img_to_array(img)
    x = np.expand_dims(x, axis=0)
    x = preprocess_input(x)
    
    sess = onnxruntime.InferenceSession(onnx_model)
    x = x if isinstance(x, list) else [x]
    feed = dict([(input.name, x[n]) for n, input in enumerate(sess.get_inputs())])
    feat = sess.run(None, feed)[0]
    
    norm_feat = feat[0] / LA.norm(feat[0])
    norm_feat = [i.item() for i in norm_feat]
    return norm_feat

Use ONNX-formatted VGG16 model to process image data:

# generate vectors with ResNet50 and VGG16 ONNX model
2vec_resnet = get_onnx_vectors("onnx_resnet50.onnx", "./pic/example.jpg")
3vec_vgg = get_onnx_vectors("onnx_vgg16.onnx", "./pic/example.jpg")

Store Vector Data

Unstructured data such as pictures cannot be directly processed by computer, but it can be converted into vectors through AI model and then analyzed by a computer. Milvus vector database is designed power massive unstructured data analysis. It can store vector data and perform near real-time analysis. First, create a collection of the corresponding model in Milvus, and then insert the image vectors.

from milvus import *

# create collections in Milvus
milvus.create_collection(resnet_collection_param)
milvus.create_collection(vgg_collection_param)

# insert data to Milvus and return ids
status, resnet_ids = milvus.insert(resnet_collection_name, resnet_vectors)
status, vgg_ids = milvus.insert(vgg_collection_name, vgg_vectors)

After inserting the data successfully, Milvus will return the ID corresponding to the vector, and then we can find the picture by ID. Since Milvus 1.1 used in this case does not support scalar filtering (which Milvus 2.0 now supports), Redis is used to store the vector ID and the key-value of the image path.

import redis
def img_ids_to_redis(img_directory, res_ids):
  for img, ids in zip(images, res_ids):
    redis.set(ids, img)

Search for Similar Images

After storing the data, we can retrieve the vector. Milvus supports multiple distance calculation methods, including Euclidean, inner product and Hamming distance. The image similarity search in this article adopts the Euclidean distance calculation between the vectors in Milvus, returns the similar vector ID, and then finds the image corresponding to the ID in Redis.

# search in Milvus and return the similarly results with ids
def search_in_milvus(collection_name, search_vector):
    status, results = milvus.search(collection_name, TOP_K, [search_vector])
    print(status)
    re_ids = [x.id for x in results[0]]
    re_distance = [x.distance for x in results[0]]
    return re_ids, re_distance
    
# get the images according the result ids
def get_sim_imgs(collection_name, search_vector):
    ids, distance = search_in_milvus(collection_name, search_vector)
    img = [red.get(i).decode("utf-8") for i in ids]
    return ids, distance, img

Taking the VGG16 and ResNet50 models as examples, this article shows processing multiple models through ONNX and combining multiple models with Milvus for similar vector retrieval to get similar images. The above two models are based on the Keras framework, which can quickly extract feature vectors. It can be seen from the Notebook that although the results of Milvus’s search for pictures on the COCO dataset based on these two models are similar, their Euclidean distances are not the same. You can also try to compare the search results of the two models using other datasets.

Milvus is a high-performance, highly available vector database that can be used to process feature vectors generated from massive unstructured data. For more solutions, you can refer to Milvus bootcamp.

References

1. https://github.com/onnx/onnx
2. https://onnx.ai/
3. https://milvus.io/cn/
4. https://github.com/milvus-io/bootcamp

About author

Shiyu Chen, a data engineer at Zilliz, graduated from Xidian University with a degree in Computer Science. Since joining Zilliz, she has been exploring solutions for Milvus in various fields, such as audio and video analysis, molecule formula retrieval, etc., which has greatly enriched the application scenarios of the community. She is currently exploring more interesting solutions. In her spare time, she loves sports and reading.