notebooks/community/vision/image_warehouse_sdk.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FX2wUzd3gjTc"
      },
      "outputs": [],
      "source": [
        "# Copyright 2023 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JAPoU8Sm5E6e"
      },
      "source": [
        "# Image Warehouse SDK demo\n",
        "\n",
        "<table align=\"left\">\n",
        "\n",
        "  <td>\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/vision/image_warehouse_sdk.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/vision/image_warehouse_sdk.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
        "      View on GitHub\n",
        "    </a>\n",
        "  </td>                                                                                         \n",
        "  <td>\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/vision/image_warehouse_sdk.ipynb\">\n",
        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
        "      Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jI8d5ytBhfuw"
      },
      "source": [
        "**_NOTE_**: This notebook has been tested in the following environment:\n",
        "\n",
        "* Python version = 3.10"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ENEut1m0h3Uo"
      },
      "source": [
        "## Overview\n",
        "\n",
        "Learn how to build a [Image Warehouse](https://cloud.google.com/vision-ai/docs) step by step by using SDK."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gMF4pbEIuZS2"
      },
      "source": [
        "### Objective\n",
        "The objective is to demostrate how to use Image Warehouse for image data ingestion and perform similarity search given text query or image as inputs. It contains examples using critical Warehouse APIs and the pipeline to perform E2E data ingestion and search journey. The colab builds a Warehouse Corpus with thousands of images ingested, analyzed and indexed; and an Index Endpoint to perform search over the images. The CUJ is as following:\n",
        "\n",
        "* Create Corpus\n",
        "* Create Data Schema\n",
        "* Import Assets\n",
        "* Analyze Corpus\n",
        "* Create Index\n",
        "* Create Index Endpoint\n",
        "* Deploy Index\n",
        "* Perform Search\n",
        "* Cleanup\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NzdLmFrajntF"
      },
      "source": [
        "### Dataset\n",
        "The dataset used in this demo is publicly accessible at [gs://cloud-samples-data/ai-platform/flowers](https://pantheon.corp.google.com/storage/browser/cloud-samples-data/ai-platform/flowers). It contains 3670 images of five kinds of flowers.\n",
        "\n",
        "The metadata file are publicly accessible at [gs://cloud-samples-data/vertex-ai-vision/warehouse/demo.jsonl](https://pantheon.corp.google.com/storage/browser/_details/cloud-samples-data/vertex-ai-vision/warehouse/demo.jsonl). It contains the annotations for each image file. The colab scans the metadata file and import the images and annotations into the warehouse."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yQV4SMLUqchE"
      },
      "source": [
        "### Costs\n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "Vertex AI Vision ([Pricing](https://cloud.google.com/vision-ai/pricing))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DQ7438Mugipn"
      },
      "source": [
        "## Installation\n",
        "\n",
        "Install the following packages required to execute this notebook.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "n_uWAql-Rnvv"
      },
      "outputs": [],
      "source": [
        "!gsutil cp gs://visionai-artifacts/visionai-0.0.6-py3-none-any.whl .\n",
        "!pip install visionai-0.0.6-py3-none-any.whl --force-reinstall"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "58707a750154"
      },
      "source": [
        "### Colab only: Uncomment the following cell to restart the kernel."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f200f10a1da3"
      },
      "outputs": [],
      "source": [
        "# import IPython\n",
        "\n",
        "# app = IPython.Application.instance()\n",
        "# app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qeRSrNTGZVBJ"
      },
      "source": [
        "## Before you begin\n",
        "\n",
        "### Set up your Google Cloud project\n",
        "\n",
        "**The following steps are required, regardless of your notebook environment.**\n",
        "\n",
        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
        "\n",
        "2. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "z3i7BDALZqeQ"
      },
      "source": [
        "#### Set your project ID\n",
        "\n",
        "**If you don't know your project ID**, try the following:\n",
        "* Run `gcloud config list`.\n",
        "* Run `gcloud projects list`.\n",
        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mOkFz5v6Z5mL"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
        "\n",
        "# Set the project id\n",
        "! gcloud config set project {PROJECT_ID}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sBCra4QMA2wR"
      },
      "source": [
        "### Authenticate your Google Cloud account\n",
        "\n",
        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "74ccc9e52986"
      },
      "source": [
        "**1. Vertex AI Workbench**\n",
        "* Do nothing as you are already authenticated."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "de775a3773ba"
      },
      "source": [
        "**2. Local JupyterLab instance, uncomment and run:**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "254614fa0c46"
      },
      "outputs": [],
      "source": [
        "# ! gcloud auth login"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ef21552ccea8"
      },
      "source": [
        "**3. Colab, uncomment and run:**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "603adbbf0532"
      },
      "outputs": [],
      "source": [
        "# from google.colab import auth\n",
        "# auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ixfp-_Ne09EK"
      },
      "source": [
        "### Set Up Other Constants"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_6ePddpJiGIb"
      },
      "outputs": [],
      "source": [
        "PROJECT_NUMBER_STR = !gcloud projects describe $PROJECT_ID --format=\"value(projectNumber)\"\n",
        "PROJECT_NUMBER = int(PROJECT_NUMBER_STR[0])\n",
        "\n",
        "# Only us-central1 is supported.\n",
        "REGION = \"us-central1\"\n",
        "\n",
        "CORPUS_DISPLAY_NAME = \"iwh demo corpus\"  # @param {type: \"string\"}\n",
        "CORPUS_DESCRIPTION = \"iwh demo corpus\"  # @param {type: \"string\"}\n",
        "\n",
        "# External users can only access PROD environment.\n",
        "ENV = \"PROD\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VxXIz2fhpbe3"
      },
      "source": [
        "### Enable API"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "shwQ0FmCBzim"
      },
      "outputs": [],
      "source": [
        "!gcloud services enable \"visionai.googleapis.com\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "t_cXC20TXi7w"
      },
      "source": [
        "### Import Libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WWmMlMTXHyCt"
      },
      "outputs": [],
      "source": [
        "import math\n",
        "import time\n",
        "\n",
        "import ipywidgets\n",
        "import requests\n",
        "from IPython.display import display\n",
        "from ipywidgets import GridspecLayout\n",
        "from visionai.python.gapic.visionai import visionai_v1\n",
        "from visionai.python.net import channel"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3oNEft7H-LFW"
      },
      "source": [
        "## Create a Warehouse client"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "17R9weLH-PSo"
      },
      "outputs": [],
      "source": [
        "warehouse_endpoint = channel.get_warehouse_service_endpoint(channel.Environment[ENV])\n",
        "warehouse_client = visionai_v1.WarehouseClient(\n",
        "    client_options={\"api_endpoint\": warehouse_endpoint}\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "P54JllW1VRZQ"
      },
      "source": [
        "## Create a Corpus"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BRFymrurNUAp"
      },
      "outputs": [],
      "source": [
        "# Set CORPUS_NAME to empty string to create new corpus\n",
        "CORPUS_NAME = \"\"  # @param {type: \"string\"}\n",
        "\n",
        "if CORPUS_NAME == \"\":\n",
        "    search_capability = visionai_v1.SearchCapability(\n",
        "        type_=visionai_v1.SearchCapability.Type.EMBEDDING_SEARCH\n",
        "    )\n",
        "    operation = warehouse_client.create_corpus(\n",
        "        visionai_v1.CreateCorpusRequest(\n",
        "            parent=f\"projects/{PROJECT_NUMBER}/locations/{REGION}\",\n",
        "            corpus=visionai_v1.Corpus(\n",
        "                display_name=CORPUS_DISPLAY_NAME,\n",
        "                description=CORPUS_DESCRIPTION,\n",
        "                type_=visionai_v1.Corpus.Type.IMAGE,\n",
        "                search_capability_setting=visionai_v1.SearchCapabilitySetting(\n",
        "                    search_capabilities=[search_capability]\n",
        "                ),\n",
        "            ),\n",
        "        )\n",
        "    )\n",
        "    print(\"Wait for corpus operation:\", operation.operation)\n",
        "\n",
        "    print(\"Created corpus \", operation.result(timeout=7200))\n",
        "    corpus_name = operation.result().name\n",
        "    print(\"Corpus created:\", corpus_name)\n",
        "else:\n",
        "    corpus_name = CORPUS_NAME\n",
        "    print(\"Corpus: \", corpus_name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZvtPdSmwVc9_"
      },
      "source": [
        "## Create DataSchema"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yrNSA_JXnzCj"
      },
      "outputs": [],
      "source": [
        "# Set SCHEMA_NAME_* to empty strings to create new schemas.\n",
        "SCHEMA_NAME_WIDTH = \"\"  # @param {type: \"string\"}\n",
        "SCHEMA_NAME_HEIGHT = \"\"  # @param {type: \"string\"}\n",
        "SCHEMA_NAME_ASPECT = \"\"  # @param {type: \"string\"}\n",
        "SCHEMA_NAME_CREATOR = \"\"  # @param {type: \"string\"}\n",
        "\n",
        "if SCHEMA_NAME_WIDTH == \"\":\n",
        "    schema_width = warehouse_client.create_data_schema(\n",
        "        visionai_v1.CreateDataSchemaRequest(\n",
        "            parent=corpus_name,\n",
        "            data_schema=visionai_v1.DataSchema(\n",
        "                key=\"width\",\n",
        "                schema_details=visionai_v1.DataSchemaDetails(\n",
        "                    type_=visionai_v1.DataSchemaDetails.DataType.STRING,\n",
        "                    granularity=visionai_v1.DataSchemaDetails.Granularity.GRANULARITY_ASSET_LEVEL,\n",
        "                    search_strategy=visionai_v1.DataSchemaDetails.SearchStrategy(\n",
        "                        search_strategy_type=visionai_v1.DataSchemaDetails.SearchStrategy.SearchStrategyType.EXACT_SEARCH\n",
        "                    ),\n",
        "                ),\n",
        "            ),\n",
        "        )\n",
        "    )\n",
        "    print(schema_width)\n",
        "    schema_name_width = schema_width.name\n",
        "else:\n",
        "    schema_name_width = SCHEMA_NAME_WIDTH\n",
        "\n",
        "if SCHEMA_NAME_HEIGHT == \"\":\n",
        "    schema_height = warehouse_client.create_data_schema(\n",
        "        visionai_v1.CreateDataSchemaRequest(\n",
        "            parent=corpus_name,\n",
        "            data_schema=visionai_v1.DataSchema(\n",
        "                key=\"height\",\n",
        "                schema_details=visionai_v1.DataSchemaDetails(\n",
        "                    type_=visionai_v1.DataSchemaDetails.DataType.STRING,\n",
        "                    granularity=visionai_v1.DataSchemaDetails.Granularity.GRANULARITY_ASSET_LEVEL,\n",
        "                    search_strategy=visionai_v1.DataSchemaDetails.SearchStrategy(\n",
        "                        search_strategy_type=visionai_v1.DataSchemaDetails.SearchStrategy.SearchStrategyType.EXACT_SEARCH\n",
        "                    ),\n",
        "                ),\n",
        "            ),\n",
        "        )\n",
        "    )\n",
        "    print(schema_height)\n",
        "    schema_name_height = schema_height.name\n",
        "else:\n",
        "    schema_name_height = SCHEMA_NAME_HEIGHT\n",
        "\n",
        "if SCHEMA_NAME_ASPECT == \"\":\n",
        "    schema_aspect = warehouse_client.create_data_schema(\n",
        "        visionai_v1.CreateDataSchemaRequest(\n",
        "            parent=corpus_name,\n",
        "            data_schema=visionai_v1.DataSchema(\n",
        "                key=\"aspect-ratio\",\n",
        "                schema_details=visionai_v1.DataSchemaDetails(\n",
        "                    type_=visionai_v1.DataSchemaDetails.DataType.STRING,\n",
        "                    granularity=visionai_v1.DataSchemaDetails.Granularity.GRANULARITY_ASSET_LEVEL,\n",
        "                    search_strategy=visionai_v1.DataSchemaDetails.SearchStrategy(\n",
        "                        search_strategy_type=visionai_v1.DataSchemaDetails.SearchStrategy.SearchStrategyType.EXACT_SEARCH\n",
        "                    ),\n",
        "                ),\n",
        "            ),\n",
        "        )\n",
        "    )\n",
        "    print(schema_aspect)\n",
        "    schema_name_aspect = schema_aspect.name\n",
        "else:\n",
        "    schema_name_aspect = SCHEMA_NAME_ASPECT\n",
        "\n",
        "if SCHEMA_NAME_CREATOR == \"\":\n",
        "    schema_creator = warehouse_client.create_data_schema(\n",
        "        visionai_v1.CreateDataSchemaRequest(\n",
        "            parent=corpus_name,\n",
        "            data_schema=visionai_v1.DataSchema(\n",
        "                key=\"creator\",\n",
        "                schema_details=visionai_v1.DataSchemaDetails(\n",
        "                    type_=visionai_v1.DataSchemaDetails.DataType.STRING,\n",
        "                    granularity=visionai_v1.DataSchemaDetails.Granularity.GRANULARITY_ASSET_LEVEL,\n",
        "                    search_strategy=visionai_v1.DataSchemaDetails.SearchStrategy(\n",
        "                        search_strategy_type=visionai_v1.DataSchemaDetails.SearchStrategy.SearchStrategyType.EXACT_SEARCH\n",
        "                    ),\n",
        "                ),\n",
        "            ),\n",
        "        )\n",
        "    )\n",
        "    print(schema_creator)\n",
        "    schema_name_creator = schema_creator.name\n",
        "else:\n",
        "    schema_name_creator = SCHEMA_NAME_CREATOR"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Mt4Z_PPB2ulp"
      },
      "source": [
        "## Import Assets"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XSnwkgjirdZ-"
      },
      "outputs": [],
      "source": [
        "# Upload images into a gcs bucket and prepare the input gcs file.\n",
        "\n",
        "# Set IMPORT_ASSET to True to import assets.\n",
        "IMPORT_ASSET = True  # @param {type: \"boolean\"}\n",
        "INPUT_GCS_FILE = \"gs://cloud-samples-data/vertex-ai-vision/warehouse/demo.jsonl\"  # @param {type: \"string\"}\n",
        "\n",
        "if IMPORT_ASSET:\n",
        "    import_lro = warehouse_client.import_assets(\n",
        "        visionai_v1.ImportAssetsRequest(\n",
        "            parent=f\"{corpus_name}\",\n",
        "            assets_gcs_uri=f\"{INPUT_GCS_FILE}\",\n",
        "        )\n",
        "    )\n",
        "    print(\"Wait for import operation: \", import_lro.operation)\n",
        "    while not import_lro.done():\n",
        "        time.sleep(10)\n",
        "    print(\"Import operation done: \", import_lro.operation)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e9F9KejLVnNA"
      },
      "source": [
        "##  Analyze Corpus"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ya64vk-7xybh"
      },
      "outputs": [],
      "source": [
        "# Set ANALYZE_CORPUS to True to analyze all assets in the corpus\n",
        "ANALYZE_CORPUS = True  # @param {type: \"boolean\"}\n",
        "\n",
        "if ANALYZE_CORPUS:\n",
        "    analyze_lro = warehouse_client.analyze_corpus(\n",
        "        visionai_v1.AnalyzeCorpusRequest(\n",
        "            name=f\"{corpus_name}\",\n",
        "        )\n",
        "    )\n",
        "    print(\"Wait for analyze operation: \", analyze_lro.operation)\n",
        "    while not analyze_lro.done():\n",
        "        time.sleep(10)\n",
        "    print(\"Analyze operation done: \", analyze_lro.operation)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "N9ZNCVXQ5xqL"
      },
      "source": [
        "## Create and deploy Index"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SCDTWM1fV4FQ"
      },
      "source": [
        "### Create Index"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kqjwPISF2lEp"
      },
      "outputs": [],
      "source": [
        "# Set INDEX_NAME to empty string to create a new index\n",
        "INDEX_NAME = \"\"  # @param {type: \"string\"}\n",
        "\n",
        "if INDEX_NAME == \"\":\n",
        "    IMAGE_INDEX_ID = \"image-index-demo\"\n",
        "    index_lro = warehouse_client.create_index(\n",
        "        visionai_v1.CreateIndexRequest(\n",
        "            parent=corpus_name,\n",
        "            index_id=f\"{IMAGE_INDEX_ID}\",\n",
        "            index=visionai_v1.Index(\n",
        "                entire_corpus=True,\n",
        "                display_name=\"demo index\",\n",
        "                description=\"demo index\",\n",
        "            ),\n",
        "        )\n",
        "    )\n",
        "    print(\"Wait for index operation:\", index_lro.operation)\n",
        "\n",
        "    print(\"Created index \", index_lro.result(timeout=10800))\n",
        "    index_name = index_lro.result().name\n",
        "    print(\"Index created:\", index_name)\n",
        "else:\n",
        "    index_name = INDEX_NAME"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GUq7XtDqV99n"
      },
      "source": [
        "### Create Index Endpoint"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jT7OloPI-sGg"
      },
      "outputs": [],
      "source": [
        "# Set INDEX_NAME to empty string to create a new index\n",
        "ENDPOINT_NAME = \"\"  # @param {type: \"string\"}\n",
        "\n",
        "if ENDPOINT_NAME == \"\":\n",
        "    ENDPOINT_ID = \"search-endpoint-demo\"\n",
        "    endpoint_lro = warehouse_client.create_index_endpoint(\n",
        "        visionai_v1.CreateIndexEndpointRequest(\n",
        "            parent=f\"projects/{PROJECT_NUMBER}/locations/{REGION}\",\n",
        "            index_endpoint_id=f\"{ENDPOINT_ID}\",\n",
        "            index_endpoint=visionai_v1.IndexEndpoint(\n",
        "                display_name=\"demo index endpoint\",\n",
        "                description=\"demo index endpoint\",\n",
        "            ),\n",
        "        )\n",
        "    )\n",
        "    print(\"Wait for endpoint operation:\", endpoint_lro.operation)\n",
        "\n",
        "    print(\"Created endpoint \", endpoint_lro.result(timeout=7200))\n",
        "    endpoint_name = endpoint_lro.result().name\n",
        "    print(\"Endpoint created:\", endpoint_name)\n",
        "else:\n",
        "    endpoint_name = ENDPOINT_NAME"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4PUBzMdGWC3_"
      },
      "source": [
        "### Deploy Index"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OdYVtCDgBFmA"
      },
      "outputs": [],
      "source": [
        "# Set DEPLOY_INDEX to True to deploy the index to the endpoint\n",
        "DEPLOY_INDEX = True  # @param {type: \"boolean\"}\n",
        "\n",
        "if DEPLOY_INDEX:\n",
        "    deploy_lro = warehouse_client.deploy_index(\n",
        "        visionai_v1.DeployIndexRequest(\n",
        "            index_endpoint=endpoint_name,\n",
        "            deployed_index=visionai_v1.DeployedIndex(\n",
        "                index=index_name,\n",
        "            ),\n",
        "        )\n",
        "    )\n",
        "    print(\"Wait for deploy operation:\", deploy_lro.operation)\n",
        "\n",
        "    print(deploy_lro.result(timeout=7200))\n",
        "    print(\"Deployed Index: \", deploy_lro.operation)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tRrvMMeLQ7lz"
      },
      "source": [
        "## Search"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0l6nQbZ--X_B"
      },
      "source": [
        "### Util for rending images"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-WTqnBAP-a78"
      },
      "outputs": [],
      "source": [
        "def RenderImages(cols=5, image_uris=[]):\n",
        "    assert len(image_uris) > 0\n",
        "    assert cols > 0\n",
        "    rows = math.floor((len(image_uris) - 1) / cols) + 1\n",
        "    grid = GridspecLayout(rows, cols)\n",
        "    for i in range(rows):\n",
        "        for j in range(cols):\n",
        "            index = i * cols + j\n",
        "            if index >= len(image_uris):\n",
        "                break\n",
        "            grid[i, j] = ipywidgets.Image(\n",
        "                value=requests.get(image_uris[index]).content, width=200\n",
        "            )\n",
        "    display(grid)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vSTndOjsWNb_"
      },
      "source": [
        "### Search by text"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "caBg19gEQ96i"
      },
      "outputs": [],
      "source": [
        "MAX_RESULTS = 10  # @param {type: \"integer\"} Set to 0 to allow all results.\n",
        "QUERY = \"multiple purple tulips\"  # @param {type: \"string\"}\n",
        "\n",
        "results = warehouse_client.search_index_endpoint(\n",
        "    visionai_v1.SearchIndexEndpointRequest(\n",
        "        index_endpoint=endpoint_name,\n",
        "        text_query=QUERY,\n",
        "    ),\n",
        ")\n",
        "\n",
        "results_cnt = 0\n",
        "asset_names = []\n",
        "for r in results:\n",
        "    asset_names.append(r.asset)\n",
        "    results_cnt += 1\n",
        "    if results_cnt >= MAX_RESULTS:\n",
        "        break\n",
        "\n",
        "uris = list(\n",
        "    map(\n",
        "        lambda asset_name: warehouse_client.generate_retrieval_url(\n",
        "            visionai_v1.GenerateRetrievalUrlRequest(\n",
        "                name=asset_name,\n",
        "            )\n",
        "        ).signed_uri,\n",
        "        asset_names,\n",
        "    )\n",
        ")\n",
        "\n",
        "RenderImages(image_uris=uris)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4T_UnnT6WQAH"
      },
      "source": [
        "### Search by image"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "s4wDqMEFDK8G"
      },
      "outputs": [],
      "source": [
        "IMAGE_GCS_FILE = \"gs://cloud-samples-data/ai-platform/flowers/roses/14312910041_b747240d56_n.jpg\"  # @#param {type: \"string\"} example: gs://iwh_fishfood/sample-image.jpg\n",
        "MAX_RESULTS = 10  # @#param {type: \"integer\"} Set to 0 to allow all results.\n",
        "IMAGE_FILE = \"/tmp/sample-image.jpg\"\n",
        "!gsutil cp $IMAGE_GCS_FILE $IMAGE_FILE\n",
        "\n",
        "with open(IMAGE_FILE, \"rb\") as f:\n",
        "    image_content = f.read()\n",
        "grid = GridspecLayout(1, 1)\n",
        "grid[0, 0] = ipywidgets.Image(value=image_content, width=200)\n",
        "\n",
        "print(\"Query image:\")\n",
        "display(grid)\n",
        "\n",
        "results = warehouse_client.search_index_endpoint(\n",
        "    visionai_v1.SearchIndexEndpointRequest(\n",
        "        index_endpoint=endpoint_name,\n",
        "        image_query=visionai_v1.ImageQuery(\n",
        "            input_image=image_content,\n",
        "        ),\n",
        "    ),\n",
        ")\n",
        "\n",
        "results_cnt = 0\n",
        "asset_names = []\n",
        "for r in results:\n",
        "    asset_names.append(r.asset)\n",
        "    results_cnt += 1\n",
        "    if results_cnt >= MAX_RESULTS:\n",
        "        break\n",
        "\n",
        "uris = list(\n",
        "    map(\n",
        "        lambda asset_name: warehouse_client.generate_retrieval_url(\n",
        "            visionai_v1.GenerateRetrievalUrlRequest(\n",
        "                name=asset_name,\n",
        "            )\n",
        "        ).signed_uri,\n",
        "        asset_names,\n",
        "    )\n",
        ")\n",
        "\n",
        "print(\"Search results:\")\n",
        "RenderImages(image_uris=uris)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uzkQSEOMbB5B"
      },
      "source": [
        "## Cleaning up"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9IpVtfcywIl9"
      },
      "outputs": [],
      "source": [
        "CLEAN_UP = False  # @param {type: \"boolean\"}\n",
        "if CLEAN_UP:\n",
        "    undeploy_lro = warehouse_client.undeploy_index(\n",
        "        visionai_v1.UndeployIndexRequest(\n",
        "            index_endpoint=endpoint_name,\n",
        "        )\n",
        "    )\n",
        "    print(\"Wait for undeploy operation:\", undeploy_lro.operation)\n",
        "\n",
        "    print(undeploy_lro.result(timeout=7200))\n",
        "\n",
        "    delete_index_lro = warehouse_client.delete_index(\n",
        "        visionai_v1.DeleteIndexRequest(\n",
        "            name=index_name,\n",
        "        )\n",
        "    )\n",
        "    print(\"Wait for delete operation:\", delete_index_lro.operation)\n",
        "\n",
        "    delete_endpoint_lro = warehouse_client.delete_index_endpoint(\n",
        "        visionai_v1.DeleteIndexEndpointRequest(\n",
        "            name=endpoint_name,\n",
        "        )\n",
        "    )\n",
        "    print(\"Wait for delete operation:\", delete_endpoint_lro.operation)\n",
        "\n",
        "    while True:\n",
        "        assets = warehouse_client.list_assets(\n",
        "            visionai_v1.ListAssetsRequest(\n",
        "                parent=corpus_name,\n",
        "                page_size=1000,\n",
        "            )\n",
        "        )\n",
        "        deletion_cnt = 0\n",
        "        for a in assets:\n",
        "            deletion_cnt += 1\n",
        "            print(\"Deleting asset:\", a.name)\n",
        "            warehouse_client.delete_asset(\n",
        "                visionai_v1.DeleteAssetRequest(\n",
        "                    name=a.name,\n",
        "                )\n",
        "            )\n",
        "            if deletion_cnt == 1000:\n",
        "                break\n",
        "        if deletion_cnt < 1000:\n",
        "            break\n",
        "\n",
        "    warehouse_client.delete_corpus(\n",
        "        visionai_v1.DeleteCorpusRequest(\n",
        "            name=corpus_name,\n",
        "        )\n",
        "    )"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "image_warehouse_sdk.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}