etc/onnx_and_milvus/milvus_with_different_models.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a76668e5",
   "metadata": {},
   "source": [
    "# Run Image Similarity Search with Milvus on Different Models\n",
    "\n",
    "In this example we will run two models to generate image feature vectors and do similarity search with [Milvus](https://milvus.io/). At the same time, we will use [ONNX](https://github.com/onnx/onnx) to manage the different models(Vgg16 and Resnet50)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "227e709d",
   "metadata": {},
   "source": [
    "## Data\n",
    "We use the [COCO animals dataset](http://cs231n.stanford.edu/coco-animals.zip), which is a smaller subset of the COCO dataset, which has 800 training images and 200 test images of 8 classes of animals: bear, bird, cat, dog, giraffe, horse, sheep, and zebra.\n",
    "\n",
    "And we extract 200 test image from it, you can download the zip file on [google drive](https://drive.google.com/file/d/1Nz67USZbz3bloUcf6PFrt7SZ7p6oMSS_/view?usp=sharing)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d186e2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# unzip the coco-images\n",
    "! unzip coco-images.zip"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b6cc2b4",
   "metadata": {},
   "source": [
    "## Requirement \n",
    "| Packages |  Servers |\n",
    "| --------------- | -------------- |\n",
    "| pymilvus_orm2.0.0rc1        | Milvus-2.0rc1   |\n",
    "| redis | Redis          |\n",
    "| onnx |       |\n",
    "| onnxruntime |       |\n",
    "| tensorflow |       |\n",
    "| tf2onnx |       |\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5139211d",
   "metadata": {},
   "source": [
    "## Up and running\n",
    "\n",
    "\n",
    "#### 1. Start Milvus server\n",
    "\n",
    "```bash\n",
    "wget https://raw.githubusercontent.com/milvus-io/milvus/master/deployments/docker/standalone/docker-compose.yml -O docker-compose.yml\n",
    "docker-compose up -d\n",
    "```\n",
    "\n",
    "This demo uses Milvus 2.0. Refer to the [Install Milvus](https://milvus.io/docs/v2.0.0/install_standalone-docker.md) for how to install Milvus docker. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73cf8e67",
   "metadata": {},
   "source": [
    "#### 2. Start Redis\n",
    "Start [Redis](https://hub.docker.com/_/redis/) with Docker container, Redis is used to store image paths and the corresponding Milvus ids."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fd7940f",
   "metadata": {},
   "outputs": [],
   "source": [
    "! docker run -d -p 6379:6379 redis"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d2b9c2d",
   "metadata": {},
   "source": [
    "#### 3. Install python packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "acdcac98",
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install -r requirements.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00fd7b2d",
   "metadata": {},
   "source": [
    "## Code Overview\n",
    "Once the environment is ready, we can start running the code. Let's start a similar image search together."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8f97b30",
   "metadata": {},
   "source": [
    "#### 1. Load ResNet50 model and covert it to ONNX\n",
    "The first step is to extract the feature vector of the image, as mentioned earlier we will run it using two models, one is the ResNet50 model based Keras and the other is the VGG16 model.\n",
    "\n",
    "As you will see in the code below, VGG16 is a pre-trained model in ONNX format, so we can convert ResNet50 to ONNX format as well to facilitate uniform use later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df7841ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "import onnxruntime\n",
    "import keras2onnx\n",
    "import numpy as np\n",
    "from numpy import linalg as LA\n",
    "from keras.applications.resnet50 import ResNet50\n",
    "from keras.preprocessing import image\n",
    "from keras.applications.resnet50 import preprocess_input\n",
    "import tensorflow as tf\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12deb1ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load keras-resnet50 model and save as a floder\n",
    "model_resnet50 = ResNet50(include_top=False, pooling='max', weights='imagenet')\n",
    "tf.saved_model.save(model_resnet50, \"keras_resnet50_model\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c1e4fa9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# convert resnet50 model to onnx\n",
    "! python -m tf2onnx.convert --saved-model \"keras_resnet50_model\" --output \"onnx_resnet50.onnx\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a086dd3",
   "metadata": {},
   "source": [
    "#### 2. Get the image embeddings with onnx model(ResNet50&VGG16)\n",
    "Next we will use the two ONNX format models to generate feature vectors. And it can be seen in the following code, both the ResNet50 model and VGG16 generate different vectors, and they do not even have the same dimensionality, Resnet50 is 2048 dimensions and vgg16 is 512 dimensions.\n",
    "\n",
    "You can download the `onnx_vgg16.onnx` file om google drive: https://drive.google.com/file/d/1QRVlLD9qgbY2K8ezM83BYyvnQZVCKuTO/view?usp=sharing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fbbfa0c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# get the image vectors with onnx model\n",
    "def get_onnx_vectors(onnx_model, img_path):\n",
    "    img = image.load_img(img_path, target_size=(224, 224))\n",
    "    x = image.img_to_array(img)\n",
    "    x = np.expand_dims(x, axis=0)\n",
    "    x = preprocess_input(x)\n",
    "    \n",
    "    sess = onnxruntime.InferenceSession(onnx_model)\n",
    "    x = x if isinstance(x, list) else [x]\n",
    "    feed = dict([(input.name, x[n]) for n, input in enumerate(sess.get_inputs())])\n",
    "    feat = sess.run(None, feed)[0]\n",
    "    \n",
    "    norm_feat = feat[0] / LA.norm(feat[0])\n",
    "    norm_feat = [i.item() for i in norm_feat]\n",
    "    return norm_feat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd94ce41",
   "metadata": {},
   "outputs": [],
   "source": [
    "# try to generate vectors with ResNet50 and prepared VGG16 ONNX model\n",
    "vec_resnet = get_onnx_vectors(\"onnx_resnet50.onnx\", \"./pic/example.jpg\")\n",
    "vec_vgg = get_onnx_vectors(\"onnx_vgg16.onnx\", \"./pic/example.jpg\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32f2b26d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# show the dimension of different models\n",
    "print(\"The dimensions of ResNet50: \", len(vec_resnet))\n",
    "print(\"The dimensions of VGG16: \", len(vec_vgg))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0efe5d18",
   "metadata": {},
   "outputs": [],
   "source": [
    "# generate the vectors batchs with the coco-images dataset\n",
    "def read_all_images(onnx_model, img_directory):\n",
    "    images = os.listdir(img_directory)\n",
    "    images.sort()\n",
    "    vectors = []\n",
    "    for image in images:\n",
    "        vec = get_onnx_vectors(onnx_model, img_directory + '/' + image)\n",
    "        vectors.append(vec)\n",
    "    return vectors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb9b3338",
   "metadata": {},
   "outputs": [],
   "source": [
    "# it will take about one minute\n",
    "IMG_DICT = 'coco-images'\n",
    "resnet_vectors = read_all_images(\"onnx_resnet50.onnx\", IMG_DICT)\n",
    "print(len(resnet_vectors))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "751f7877",
   "metadata": {},
   "outputs": [],
   "source": [
    "# it will take about one minute\n",
    "vgg_vectors = read_all_images(\"onnx_vgg16.onnx\", IMG_DICT)\n",
    "print(len(vgg_vectors))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e97e855",
   "metadata": {},
   "source": [
    "#### 3. Create collection and insert the vectors to Milvus\n",
    "In the previous section we generate the feature vectors by the model, next we can process the vectors.\n",
    "\n",
    "Milvus is an open source vector search engine, we can store and retrieve the vectors in Milvus, we first have to create the collection and insert the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b242e2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# connect to Milvus server\n",
    "from pymilvus_orm import *\n",
    "connections.connect(host=\"127.0.0.1\", port=19530)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85db1a44",
   "metadata": {},
   "outputs": [],
   "source": [
    "resnet_collection_name = \"resnet_img_search\"\n",
    "resnet_VECTOR_DIMENSION = 2048\n",
    "resnet_fields = [\n",
    "    schema.FieldSchema(name=\"id\", dtype=DataType.INT64, is_primary=True, auto_id=True),\n",
    "    schema.FieldSchema(name=\"vector\", dtype=DataType.FLOAT_VECTOR, dim=resnet_VECTOR_DIMENSION)\n",
    "]\n",
    "\n",
    "vgg_collection_name = \"vgg_img_search\"\n",
    "vgg_VECTOR_DIMENSION = 512\n",
    "vgg_fields = [\n",
    "    schema.FieldSchema(name=\"id\", dtype=DataType.INT64, is_primary=True, auto_id=True),\n",
    "    schema.FieldSchema(name=\"vector\", dtype=DataType.FLOAT_VECTOR, dim=vgg_VECTOR_DIMENSION)\n",
    "]\n",
    "\n",
    "TOP_K = 5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a0144a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Creat a collection\n",
    "\n",
    "resnet_schema = schema.CollectionSchema(fields=resnet_fields, description=\"Resnet test collection\")\n",
    "resnet_collection = Collection(name=resnet_collection_name, schema=resnet_schema)\n",
    "\n",
    "vgg_schema = schema.CollectionSchema(fields=vgg_fields, description=\"Vgg test collection\")\n",
    "vgg_collection = Collection(name=vgg_collection_name, schema=vgg_schema)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "654c1cab",
   "metadata": {},
   "outputs": [],
   "source": [
    "# insert data to Milvus and return ids\n",
    "resnet_mr = resnet_collection.insert([resnet_vectors])\n",
    "resnet_collection.load()\n",
    "resnet_ids = resnet_mr.primary_keys\n",
    "print(\"Insert resnet vector to Milvus and the len:\", len(resnet_ids))\n",
    "\n",
    "vgg_mr = vgg_collection.insert([vgg_vectors])\n",
    "vgg_collection.load()\n",
    "vgg_ids = vgg_mr.primary_keys\n",
    "print(\"Insert vgg vector to Milvus and the len:\", len(vgg_ids))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52b15b8c",
   "metadata": {},
   "source": [
    "#### 4. Insert the image and ids to Reids\n",
    "After the vectors of the images are stored in Milvus, the corresponding ids are generated, and the next step is to store the correspondence between the ids and the image paths in Redis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e4003a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# connect to Redis\n",
    "import redis\n",
    "red = redis.Redis(host = '127.0.0.1', port=6379, db=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f92914f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def img_ids_to_redis(img_directory, res_ids):\n",
    "    images = os.listdir(img_directory)\n",
    "    images.sort()\n",
    "    for img, ids in zip(images, res_ids):\n",
    "        red.set(ids, img)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76705480",
   "metadata": {},
   "outputs": [],
   "source": [
    "img_ids_to_redis(IMG_DICT, resnet_ids)\n",
    "img_ids_to_redis(IMG_DICT, vgg_ids)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c15637a",
   "metadata": {},
   "source": [
    "#### 5. Search the similarity images\n",
    "Then we can retrieve an image in Milvus, and Milvus will return the most similar image id, and then get the image according to the id in Redis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "946171af",
   "metadata": {},
   "outputs": [],
   "source": [
    "# search in Milvus and return the similarly results with ids\n",
    "def search_in_milvus(collection, search_vector):\n",
    "    search_params = {\"metric_type\": \"L2\", \"params\": {\"nprobe\": 32}}\n",
    "    results = collection.search([search_vector], \"vector\", param=search_params, limit=10, expr=None)\n",
    "    re_ids = [x.id for x in results[0]]\n",
    "    re_distance = [x.distance for x in results[0]]\n",
    "    return re_ids, re_distance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "732a10b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# get the images according the result ids\n",
    "def get_sim_imgs(collection, search_vector):\n",
    "    ids, distance = search_in_milvus(collection, search_vector)\n",
    "    img = [red.get(i).decode(\"utf-8\") for i in ids]\n",
    "    return ids, distance, img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8a96490",
   "metadata": {},
   "outputs": [],
   "source": [
    "resnet_ids, resnet_distance, resnet_img =  get_sim_imgs(resnet_collection, vec_resnet)\n",
    "vgg_ids, vgg_distance, vgg_img =  get_sim_imgs(vgg_collection, vec_vgg)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7cd2740e",
   "metadata": {},
   "source": [
    "#### 6. Show the results\n",
    "Finally we get the most similar image results(id, distance and image) to the target image, and we can see together how the different models retrieve the results.\n",
    "\n",
    "We can see that the first image of the retrieval result is our target image (because the distance is 0), and the first 5 results look the same but the distance is not the same. It can be seen that the retrieval effect of VGG16 and ResNet50 is about the same, but of course you can use more and larger data to retrieve, and the effect may be different."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cef337bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import Image, display\n",
    "def show_results(img, ids=None, distance=None):\n",
    "    for i in range(len(img)):\n",
    "        print(ids[i], distance[i])\n",
    "        x = Image(filename=IMG_DICT+'/'+resnet_img[i]) \n",
    "        display(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bfce5a14",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_results(resnet_img, resnet_ids, resnet_distance)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0925e3f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_results(vgg_img, vgg_ids, vgg_distance)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "milvus20",
   "language": "python",
   "name": "milvus2_0"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}