notebooks/official/explainable_ai/xai_image_classification_feature_attributions.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ur8xi4C7S06n"
      },
      "outputs": [],
      "source": [
        "# Copyright 2022 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": [
        "# Explaining image classification with Vertex Explainable AI\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/explainable_ai/xai_image_classification_feature_attributions.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fofficial%2Fexplainable_ai%2fxai_image_classification_feature_attributions.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
        "  </td>    \n",
        "  <td style=\"text-align: center\">\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/official/explainable_ai/xai_image_classification_feature_attributions.ipynb\">\n",
        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/explainable_ai/xai_image_classification_feature_attributions.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "## Overview\n",
        "\n",
        "Vertex Explainable AI offers feature-based and example-based explanations to provide better understanding of model decision making. For feature-based explanations, Vertex Explainable AI integrates feature attributions into Vertex AI. Feature attributions indicate how much each feature in your model contributed to the predictions for each given instance. For an image classification model, when you request explanations, you get the predicted class along with an overlay for the image, showing which areas in the image contributed most strongly to the resulting prediction.\n",
        "\n",
        "To use Vertex Explainable AI on a pre-trained or custom-trained model, you must configure certain options when you create the model resource that you plan to request explanations from, when you deploy the model, or when you submit a batch explanation job. This tutorial demonstrates how to configure these options, and get and visualize explanations for online and batch predictions.\n",
        "\n",
        "Learn more about [Vertex Explainable AI](https://cloud.google.com/vertex-ai/docs/explainable-ai/overview) and [Vertex AI Prediction](https://cloud.google.com/vertex-ai/docs/predictions/get-predictions)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d975e698c9a4"
      },
      "source": [
        "### Objective\n",
        "\n",
        "In this tutorial, you learn how to configure feature-based explanations on a pre-trained image classification model and make online and batch predictions with explanations.\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services and resources:\n",
        "\n",
        "- Vertex Explainable AI\n",
        "- Vertex AI Prediction\n",
        "\n",
        "The steps performed include:\n",
        "\n",
        "- Download pretrained model from TensorFlow Hub\n",
        "- Upload model for deployment\n",
        "- Deploy model for online prediction\n",
        "- Make online prediction with explanations\n",
        "- Make batch predictions with explanations\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "08d289fa873f"
      },
      "source": [
        "### Dataset\n",
        "\n",
        "In this example, you use the TensorFlow [flowers](http://download.tensorflow.org/example_images/flower_photos.tgz) dataset. The dataset contains about 3,700 photos of flowers in five sub-directories, one per class."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aed92deeb4a0"
      },
      "source": [
        "### Costs\n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "* Vertex AI\n",
        "* Cloud Storage\n",
        "\n",
        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing),\n",
        "and [Cloud Storage pricing](https://cloud.google.com/storage/pricing),\n",
        "and use the [Pricing Calculator](https://cloud.google.com/products/calculator/)\n",
        "to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "61RBz8LLbxCR"
      },
      "source": [
        "## Get started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "No17Cw5hgx12"
      },
      "source": [
        "### Install Vertex AI SDK for Python and other required packages\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2b4ef9b72d43"
      },
      "outputs": [],
      "source": [
        "# Install the packages\n",
        "! pip3 install --upgrade -q google-cloud-aiplatform \\\n",
        "                            tensorflow==2.15.1 \\\n",
        "                            tensorflow-hub \\\n",
        "                            matplotlib"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R5Xep4W9lq-Z"
      },
      "source": [
        "### Restart runtime (Colab only)\n",
        "\n",
        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XRvKdaPDTznN"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "\n",
        "    import IPython\n",
        "\n",
        "    app = IPython.Application.instance()\n",
        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SbmM4z7FOBpM"
      },
      "source": [
        "<div class=\"alert alert-block alert-warning\">\n",
        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
        "</div>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dmWOrTJ3gx13"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
        "Authenticate your environment on Google Colab.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NyKGtVQjgx13"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DF4l8DTdWgPY"
      },
      "source": [
        "### Set Google Cloud project information\n",
        "\n",
        "Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Nqwi-5ufWp_B"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
        "LOCATION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zgPO1eR3CYjk"
      },
      "source": [
        "### Create a Cloud Storage bucket\n",
        "\n",
        "Create a storage bucket to store intermediate artifacts such as datasets."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MzGDU7TWdts_"
      },
      "outputs": [],
      "source": [
        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-EcIXiGsCePi"
      },
      "source": [
        "**If your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NIq7R4HZCfIc"
      },
      "outputs": [],
      "source": [
        "! gsutil mb -l $LOCATION -p $PROJECT_ID $BUCKET_URI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "960505627ddf"
      },
      "source": [
        "### Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PyQmSRbKA8r-"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "import tensorflow_hub as hub\n",
        "from google.cloud import aiplatform"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "init_aip:mbsdk,all"
      },
      "source": [
        "### Initialize Vertex AI SDK for Python\n",
        "\n",
        "To get started using Vertex AI, you must [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "utUNuq2aARaE"
      },
      "outputs": [],
      "source": [
        "aiplatform.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SqI_MEPvWVI5"
      },
      "source": [
        "## Download pre-trained model from TensorFlow Hub\n",
        "\n",
        "Feature attribution is supported for all types of models (both AutoML and custom-trained), frameworks (TensorFlow, scikit, XGBoost), and modalities (images, text, tabular, video).\n",
        "\n",
        "For demonstration purposes, this tutorial uses an image model [Inception_v3](https://tfhub.dev/google/imagenet/inception_v3/classification/5) from the TensorFlow Hub. This model was pre-trained on the ImageNet benchmark dataset.\n",
        "First, you download the model from the TensorFlow Hub, wrap it as a Keras layer with `hub.KerasLayer` and save the model artifacts to your Cloud Storage bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3_JDDMUjFfXO"
      },
      "outputs": [],
      "source": [
        "classifier_model = \"https://tfhub.dev/google/imagenet/inception_v3/classification/5\"\n",
        "\n",
        "classifier = tf.keras.Sequential([hub.KerasLayer(classifier_model)])\n",
        "\n",
        "classifier.build([None, 224, 224, 3])\n",
        "\n",
        "MODEL_DIR = f\"{BUCKET_URI}/model\"\n",
        "classifier.save(MODEL_DIR)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ySldGR7eKdY0"
      },
      "source": [
        "## Upload model for deployment\n",
        "\n",
        "Next, you upload the model to Vertex AI Model Registry, which will create a Vertex AI Model resource for your model. Prior to uploading, you need to define a serving function to convert data to the format your model expects."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "p6wXBlaeZs3j"
      },
      "source": [
        "### Define a serving function for image data\n",
        "\n",
        "You define a serving function to convert image data to the format your model expects. When you send encoded data to Vertex AI, your serving function ensures that the data is decoded on the model server before it is passed as input to your model.\n",
        "\n",
        "To enable Vertex Explainable AI in your custom models, you need to set two additional signatures from the serving function:\n",
        "\n",
        "- `xai_preprocess`: The preprocessing function in the serving function.\n",
        "- `xai_model`: The concrete function for calling the model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qzA0aHeLFfeq"
      },
      "outputs": [],
      "source": [
        "CONCRETE_INPUT = \"numpy_inputs\"\n",
        "\n",
        "\n",
        "def _preprocess(bytes_input):\n",
        "    decoded = tf.io.decode_jpeg(bytes_input, channels=3)\n",
        "    decoded = tf.image.convert_image_dtype(decoded, tf.float32)\n",
        "    resized = tf.image.resize(decoded, size=(224, 224))\n",
        "    return resized\n",
        "\n",
        "\n",
        "@tf.function(input_signature=[tf.TensorSpec([None], tf.string)])\n",
        "def preprocess_fn(bytes_inputs):\n",
        "    decoded_images = tf.map_fn(\n",
        "        _preprocess, bytes_inputs, dtype=tf.float32, back_prop=False\n",
        "    )\n",
        "    return {\n",
        "        CONCRETE_INPUT: decoded_images\n",
        "    }  # User needs to make sure the key matches model's input\n",
        "\n",
        "\n",
        "@tf.function(input_signature=[tf.TensorSpec([None], tf.string)])\n",
        "def serving_fn(bytes_inputs):\n",
        "    images = preprocess_fn(bytes_inputs)\n",
        "    prob = m_call(**images)\n",
        "    return prob\n",
        "\n",
        "\n",
        "m_call = tf.function(classifier.call).get_concrete_function(\n",
        "    [tf.TensorSpec(shape=[None, 224, 224, 3], dtype=tf.float32, name=CONCRETE_INPUT)]\n",
        ")\n",
        "\n",
        "tf.saved_model.save(\n",
        "    classifier,\n",
        "    MODEL_DIR,\n",
        "    signatures={\n",
        "        \"serving_default\": serving_fn,\n",
        "        \"xai_preprocess\": preprocess_fn,  # Required for XAI\n",
        "        \"xai_model\": m_call,  # Required for XAI\n",
        "    },\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YOOXywwNZBN1"
      },
      "source": [
        "### Get the serving function signature\n",
        "\n",
        "You get the signatures of your model's input and output layers by reloading the model into memory, and querying it for the signatures corresponding to each layer. The input layer name of the serving function will be used later when you make a prediction request."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vVdSjoBSQaEA"
      },
      "outputs": [],
      "source": [
        "loaded = tf.saved_model.load(MODEL_DIR)\n",
        "\n",
        "serving_input = list(\n",
        "    loaded.signatures[\"serving_default\"].structured_input_signature[1].keys()\n",
        ")[0]\n",
        "print(\"Serving function input:\", serving_input)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wnefcZK1ba4b"
      },
      "source": [
        "## Configure explanation settings\n",
        "\n",
        "To use Vertex Explainable AI with a custom-trained model, you must configure explanation settings when uploading the model. These settings include:\n",
        "\n",
        "- `parameters`: The feature attribution method. Available methods include `shapley`, `ig`, `xrai`.\n",
        "- `metadata`: The model's input and output for explanation. **This field is optional for TensorFlow 2 models. If omitted, Vertex AI automatically infers the inputs and outputs from the model**. You don't need to configure this field in this tutorial.\n",
        "\n",
        "Learn more about [configuring feature-based explanations](https://cloud.google.com/vertex-ai/docs/explainable-ai/configuring-explanations-feature-based)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_wDTAJoJHx5p"
      },
      "outputs": [],
      "source": [
        "XAI = \"ig\"  # [ shapley, ig, xrai ]\n",
        "\n",
        "if XAI == \"shapley\":\n",
        "    PARAMETERS = {\"sampled_shapley_attribution\": {\"path_count\": 10}}\n",
        "elif XAI == \"ig\":\n",
        "    PARAMETERS = {\"integrated_gradients_attribution\": {\"step_count\": 50}}\n",
        "elif XAI == \"xrai\":\n",
        "    PARAMETERS = {\"xrai_attribution\": {\"step_count\": 50}}\n",
        "\n",
        "parameters = aiplatform.explain.ExplanationParameters(PARAMETERS)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aLslkKlneNyX"
      },
      "source": [
        "## Upload the model to a Vertex AI Model resource\n",
        "\n",
        "Next, upload your model to a Vertex AI Model resource with the explanation configuration. Vertex AI provides [Docker container images](https://cloud.google.com/vertex-ai/docs/predictions/pre-built-containers) that you run as pre-built containers for serving predictions and explanations from trained model artifacts."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MQ_oIG11ID1o"
      },
      "outputs": [],
      "source": [
        "MODEL_DISPLAY_NAME = \"inception-v3-model-unique\"\n",
        "DEPLOY_IMAGE = \"us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-11:latest\"\n",
        "\n",
        "model = aiplatform.Model.upload(\n",
        "    display_name=MODEL_DISPLAY_NAME,\n",
        "    artifact_uri=MODEL_DIR,\n",
        "    serving_container_image_uri=DEPLOY_IMAGE,\n",
        "    explanation_parameters=parameters,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KfL8qvYlLOo6"
      },
      "source": [
        "## Deploy model for online prediction\n",
        "\n",
        "Next, deploy your model for online prediction. You set the variable `DEPLOY_COMPUTE` to configure the machine type for the [compute resources](https://cloud.google.com/vertex-ai/docs/predictions/configure-compute) you will use for prediction."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IOkThfvqJlk6"
      },
      "outputs": [],
      "source": [
        "DEPLOY_DISPLAY_NAME = \"inception-v3-deploy-unique\"\n",
        "DEPLOY_COMPUTE = \"n1-standard-4\"\n",
        "\n",
        "endpoint = model.deploy(\n",
        "    deployed_model_display_name=DEPLOY_DISPLAY_NAME,\n",
        "    machine_type=DEPLOY_COMPUTE,\n",
        "    accelerator_type=None,\n",
        "    accelerator_count=0,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bHCXBWAYNx3h"
      },
      "source": [
        "## Make online prediction with explanations\n",
        "\n",
        "\n",
        "### Download an image dataset and labels\n",
        "In this example, you use the TensorFlow flowers dataset for the input data for predictions. The dataset contains about 3,700 photos of flowers in five sub-directories, one per class. You also fetch the ImageNet dataset labels to decode the predictions.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e3caSDUXReNo"
      },
      "outputs": [],
      "source": [
        "import pathlib\n",
        "\n",
        "data_dir = tf.keras.utils.get_file(\n",
        "    \"flower_photos\",\n",
        "    origin=\"https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz\",\n",
        "    untar=True,\n",
        ")\n",
        "\n",
        "data_dir = pathlib.Path(data_dir)\n",
        "images_files = list(data_dir.glob(\"daisy/*\"))\n",
        "\n",
        "labels_path = tf.keras.utils.get_file(\n",
        "    \"ImageNetLabels.txt\",\n",
        "    \"https://storage.googleapis.com/download.tensorflow.org/data/ImageNetLabels.txt\",\n",
        ")\n",
        "\n",
        "imagenet_labels = np.array(open(labels_path).read().splitlines())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2OcVOWLuv6TQ"
      },
      "source": [
        "### Prepare image processing functions\n",
        "\n",
        "You define some reusable functions for image processing and visualization."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DQz4T3h8iU2M"
      },
      "outputs": [],
      "source": [
        "import base64\n",
        "import io\n",
        "import json\n",
        "\n",
        "import matplotlib.image as mpimg\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "\n",
        "def encode_image_in_b64(image_file):\n",
        "    bytes = tf.io.read_file(image_file)\n",
        "    b64str = base64.b64encode(bytes.numpy()).decode(\"utf-8\")\n",
        "    return b64str\n",
        "\n",
        "\n",
        "def decode_b64_image(b64_image_str):\n",
        "    image = base64.b64decode(b64_image_str)\n",
        "    image = io.BytesIO(image)\n",
        "    image = mpimg.imread(image, format=\"JPG\")\n",
        "    return image\n",
        "\n",
        "\n",
        "def decode_numpy_image(numpy_inputs):\n",
        "    numpy_inputs_json = json.loads(str(numpy_inputs))\n",
        "    image = np.array(numpy_inputs_json)\n",
        "    return image\n",
        "\n",
        "\n",
        "def show_explanation(encoded_image, prediction, feature_attributions):\n",
        "    fig, axs = plt.subplots(nrows=1, ncols=3, figsize=(16, 4))\n",
        "\n",
        "    label_index = np.argmax(prediction)\n",
        "    class_name = imagenet_labels[label_index]\n",
        "    confidence_score = prediction[label_index]\n",
        "    axs[0].set_title(\n",
        "        \"Prediction:[\" + class_name + \"] (\" + str(round(confidence_score, 1)) + \"%)\"\n",
        "    )\n",
        "    original_image = decode_b64_image(encoded_image)\n",
        "    axs[0].imshow(original_image, interpolation=\"nearest\", aspect=\"auto\")\n",
        "    axs[0].axis(\"off\")\n",
        "\n",
        "    numpy_inputs = feature_attributions[\"numpy_inputs\"]\n",
        "    attribution_image = decode_numpy_image(numpy_inputs)\n",
        "    axs[1].set_title(\"Feature attributions\")\n",
        "    axs[1].imshow(attribution_image, interpolation=\"nearest\", aspect=\"auto\")\n",
        "    axs[1].axis(\"off\")\n",
        "\n",
        "    processed_image = attribution_image.max(axis=2)\n",
        "    axs[2].imshow(processed_image, cmap=\"coolwarm\", aspect=\"auto\")\n",
        "    axs[2].set_title(\"Feature attribution heatmap\")\n",
        "    axs[2].axis(\"off\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NaO8nWWQxuKw"
      },
      "source": [
        "### Get online explanations\n",
        "\n",
        "You send an `explain` request with encoded input image data to the `endpoint` and get predictions with explanations."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "N_BeaaCy5M1H"
      },
      "outputs": [],
      "source": [
        "TEST_IMAGE_SIZE = 2\n",
        "\n",
        "test_image_list = []\n",
        "for i in range(TEST_IMAGE_SIZE):\n",
        "    test_image_list.append(str(images_files[i]))\n",
        "\n",
        "instances_list = []\n",
        "for test_image in test_image_list:\n",
        "    b64str = encode_image_in_b64(test_image)\n",
        "    instances_list.append({serving_input: {\"b64\": b64str}})\n",
        "\n",
        "response = endpoint.explain(instances_list)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CTegf0H-1M3M"
      },
      "source": [
        "### Visualize online explanations\n",
        "\n",
        "As you request explanations on an image classification model, you get the predicted class along with an image overlay showing which pixels (integrated gradients) or regions (integrated gradients or XRAI) contributed to the prediction."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "A69Sxd3D0On0"
      },
      "outputs": [],
      "source": [
        "for i, test_image in enumerate(test_image_list):\n",
        "    encoded_image = encode_image_in_b64(test_image)\n",
        "    prediction = response.predictions[i]\n",
        "    explanation = response.explanations[i]\n",
        "    feature_attributions = dict(explanation.attributions[0].feature_attributions)\n",
        "\n",
        "    show_explanation(encoded_image, prediction, feature_attributions)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KoJiAyLOM2tr"
      },
      "source": [
        "## Make batch predictions with explanations\n",
        "\n",
        "### Create the batch input file\n",
        "\n",
        "You create a batch input file in JSONL format and store the input file in your Cloud Storage bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5LgFY93u0PVf"
      },
      "outputs": [],
      "source": [
        "gcs_input_uri = f\"{BUCKET_URI}/test_images.json\"\n",
        "\n",
        "with tf.io.gfile.GFile(gcs_input_uri, \"w\") as f:\n",
        "    for test_image in test_image_list:\n",
        "        b64str = encode_image_in_b64(test_image)\n",
        "        data = {serving_input: {\"b64\": b64str}}\n",
        "        f.write(json.dumps(data) + \"\\n\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cgxVGYmX2Onc"
      },
      "source": [
        "### Submit a batch prediction job\n",
        "\n",
        "You make a batch prediction by submitting a batch prediction job with the `generate_explanation` parameter set to `True`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gaB_dz3h0PYe"
      },
      "outputs": [],
      "source": [
        "JOB_DISPLAY_NAME = \"inception-v3-job-unique\"\n",
        "\n",
        "batch_predict_job = model.batch_predict(\n",
        "    job_display_name=JOB_DISPLAY_NAME,\n",
        "    gcs_source=gcs_input_uri,\n",
        "    gcs_destination_prefix=BUCKET_URI,\n",
        "    instances_format=\"jsonl\",\n",
        "    model_parameters=None,\n",
        "    machine_type=DEPLOY_COMPUTE,\n",
        "    generate_explanation=True,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RVxOqx-P4nma"
      },
      "source": [
        "### Get batch explanations\n",
        "\n",
        "Next, you get the explanations from the completed batch prediction job. The results are written to the Cloud Storage output bucket you specified in the batch prediction request. You call the method `iter_outputs()` to get a list of each Cloud Storage file generated with the results."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pD7kCit00Pbg"
      },
      "outputs": [],
      "source": [
        "bp_iter_outputs = batch_predict_job.iter_outputs()\n",
        "\n",
        "explanation_results = list()\n",
        "for blob in bp_iter_outputs:\n",
        "    if blob.name.split(\"/\")[-1].startswith(\"explanation.results\"):\n",
        "        explanation_results.append(blob.name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JSdO0Dx444Ir"
      },
      "source": [
        "### Visualize explanations\n",
        "\n",
        "You take one explanation result as an example and visualize the explanations. For an image classification model, you get the predicted class along with an image overlay showing which pixels (integrated gradients) or regions (integrated gradients or XRAI) contributed to the prediction."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SJpLkjc7Fb9T"
      },
      "outputs": [],
      "source": [
        "explanation_result = explanation_results[0]\n",
        "\n",
        "gfile_name = f\"{BUCKET_URI}/{explanation_result}\"\n",
        "with tf.io.gfile.GFile(name=gfile_name, mode=\"r\") as gfile:\n",
        "    result = json.loads(gfile.read())\n",
        "\n",
        "encoded_image = result[\"instance\"][\"bytes_inputs\"][\"b64\"]\n",
        "prediction = result[\"prediction\"]\n",
        "attributions = result[\"explanation\"][\"attributions\"][0]\n",
        "feature_attributions = attributions[\"featureAttributions\"]\n",
        "\n",
        "show_explanation(encoded_image, prediction, feature_attributions)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TpV-iwP9qw9c"
      },
      "source": [
        "## Cleaning up\n",
        "\n",
        "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n",
        "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
        "\n",
        "Otherwise, you can delete the individual resources you created in this tutorial."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sx_vKniMq9ZX"
      },
      "outputs": [],
      "source": [
        "# Deploy the model from endpoint\n",
        "endpoint.undeploy_all()\n",
        "# Delete the endpoint\n",
        "endpoint.delete()\n",
        "# Delete the model\n",
        "model.delete()\n",
        "# Delete the batch prediction job\n",
        "batch_predict_job.delete()\n",
        "\n",
        "# Delete Cloud Storage bucket\n",
        "delete_bucket = False  # Set True to delete your bucket\n",
        "if delete_bucket:\n",
        "    ! gsutil -m rm -r $BUCKET_URI"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "xai_image_classification_feature_attributions.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}