{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "698f53b0-b7a6-4e87-a801-1399ec7eea89",
   "metadata": {},
   "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",
   "id": "e6ecf407-9acb-4d37-872a-ca59cf0858af",
   "metadata": {},
   "source": [
    "# Multimodal content enrichment"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ca356aa-a272-4d30-9bf8-330b82d805c1",
   "metadata": {},
   "source": [
    "## Overview"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b72e097-456e-42f1-8240-b3d595378ee4",
   "metadata": {},
   "source": [
    "This notebook shows how to process ads banners images and generate enriched information classifications using a taxonomy using Bigframes and Gemini.  \n",
    "It reads ads banners images from a GCS bucket and process them using Bigframes and Gemini, leveraging BigQuery's distributed processing capabilities.  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83cd4547-0fe5-4bf3-b46b-6b1bddc4d867",
   "metadata": {},
   "source": [
    "### Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a42e4ebf",
   "metadata": {},
   "source": [
    "### Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8d4a5a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "PROJECT_ID = \"\"  # @param {type:\"string\"}\n",
    "if not PROJECT_ID:\n",
    "    PROJECT_ID = input(\"Enter project id\")\n",
    "\n",
    "REGION = \"us-central1\"  # @param {type:\"string\"}\n",
    "if not REGION:\n",
    "    REGION = input(\"Enter region (e.g. us-central1)\")\n",
    "\n",
    "# Choose the name of the resources that will be created:\n",
    "DATASET_BUCKET_NAME = \"\"  # @param {type:\"string\"}\n",
    "if not DATASET_BUCKET_NAME:\n",
    "    DATASET_BUCKET_NAME = input(\"Enter GCS bucket name (without gs://)\")\n",
    "\n",
    "DATASET_ID = \"multimodal_ads_enrichment\"  # @param {type:\"string\"}\n",
    "OBJECT_TABLE_NAME = \"ads_banners_object_table\"  # @param {type:\"string\"}\n",
    "OUTPUT_TABLE_NAME = \"ads_banners_predictions\"  # @param {type:\"string\"}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9be56017-6e95-48c4-a8d2-ee3d76586d42",
   "metadata": {},
   "source": [
    "We are using the default compute service account for the remote function. We will grant it the necessary permissions just before creating the function in the steps below. You need to ensure you have permissions to grant IAM roles in this project (e.g., Project IAM Admin)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cfd909d4",
   "metadata": {},
   "source": [
    "#### Enable required APIs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f3466bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "!gcloud services enable cloudfunctions.googleapis.com\n",
    "!gcloud services enable cloudbuild.googleapis.com"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92489bee-d8a5-4e43-9904-71c0a03c6278",
   "metadata": {},
   "source": [
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcc0cbc9-1b92-4707-8804-554568d9ead1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!pip3 install --upgrade google-cloud-bigquery google-cloud-storage google-cloud-bigquery-connection bigframes -q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76c43815",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "import io\n",
    "\n",
    "import bigframes.pandas as bpd\n",
    "import bigframes.bigquery as bbq\n",
    "bpd.options.display.progress_bar = None\n",
    "import pandas as pd\n",
    "\n",
    "from google.cloud import storage\n",
    "from IPython.display import display\n",
    "from PIL import Image\n",
    "\n",
    "import requests\n",
    "\n",
    "from tabulate import tabulate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9344a067",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "%load_ext google.cloud.bigquery"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90fc31b5",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# Note: The project option is not required in all environments.\n",
    "# On BigQuery Studio, the project ID is automatically detected.\n",
    "bpd.options.bigquery.project = PROJECT_ID\n",
    "# Note: The location option is not required.\n",
    "# It defaults to the location of the first table or query\n",
    "# passed to read_gbq(). For APIs where a location can't be\n",
    "# auto-detected, the location defaults to the \"US\" location.\n",
    "bpd.options.bigquery.location = REGION"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4735583d",
   "metadata": {},
   "source": [
    "### Copy data to your own bucket"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1df4c89d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_bucket_and_copy_data(source_path: str, destination_path: str, location: str):\n",
    "    \"\"\"\n",
    "    Creates a GCS bucket and copies all data from a source path to the new bucket\n",
    "    using a more compatible copy method.\n",
    "\n",
    "    Args:\n",
    "        source_path (str): The 'gs://' path of the source bucket.\n",
    "        destination_path (str): The 'gs://' path for the new bucket to be created.\n",
    "        location (str): The location for the new bucket (e.g., 'us-central1').\n",
    "    \"\"\"\n",
    "    # Parse bucket names from the 'gs://' paths\n",
    "    source_bucket_name = source_path.replace(\"gs://\", \"\").split(\"/\")[0]\n",
    "    destination_bucket_name = destination_path.replace(\"gs://\", \"\")\n",
    "\n",
    "    storage_client = storage.Client()\n",
    "\n",
    "    # Create the new destination bucket\n",
    "    print(f\"Creating bucket: {destination_bucket_name} in location: {location}\")\n",
    "    try:\n",
    "        destination_bucket = storage_client.create_bucket(destination_bucket_name, location=location)\n",
    "        print(f\"✅ Bucket '{destination_bucket.name}' created successfully.\")\n",
    "    except Exception as e:\n",
    "        print(f\"⚠️ Error creating bucket. It may already exist: {e}\")\n",
    "        destination_bucket = storage_client.bucket(destination_bucket_name)\n",
    "\n",
    "    # Get the source bucket\n",
    "    source_bucket = storage_client.bucket(source_bucket_name)\n",
    "\n",
    "    # Get the folder name from the source path, if any\n",
    "    source_folder = source_path.replace(f\"gs://{source_bucket_name}/\", \"\")\n",
    "    if source_folder == source_bucket_name:\n",
    "        source_folder = \"\"\n",
    "\n",
    "    # List and copy blobs from the source path to the destination\n",
    "    print(f\"\\nCopying data from '{source_path}' to '{destination_bucket_name}'...\")\n",
    "    \n",
    "    # We iterate over blobs in the source bucket, filtered by the prefix\n",
    "    blobs = source_bucket.list_blobs(prefix=source_folder)\n",
    "    \n",
    "    copied_count = 0\n",
    "    for blob in blobs:\n",
    "        # Create a new blob name for the destination, removing the source folder prefix\n",
    "        destination_blob_name = blob.name.replace(source_folder, \"\", 1).lstrip('/')\n",
    "        \n",
    "        # If the destination blob name is empty, it means we are trying to copy a folder itself, skip.\n",
    "        if not destination_blob_name:\n",
    "            continue\n",
    "\n",
    "        # Check if the blob is a folder representation (ends with /), if so, skip.\n",
    "        if blob.name.endswith('/'):\n",
    "            continue\n",
    "            \n",
    "        # Use the more compatible copy_blob() method\n",
    "        source_bucket.copy_blob(blob, destination_bucket, destination_blob_name)\n",
    "        print(f\"   ➡️ Copied '{blob.name}' to '{destination_blob_name}'\")\n",
    "        copied_count += 1\n",
    "    \n",
    "    print(f\"\\n✅ Finished copying {copied_count} file(s).\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a4aa252",
   "metadata": {},
   "outputs": [],
   "source": [
    "SOURCE_BUCKET_PATH = \"gs://dataproc-metastore-public-binaries/ads_banners_images\" # Public dataset\n",
    "DATASET_PATH = f\"gs://{DATASET_BUCKET_NAME}\"\n",
    "DATASET_PATH_LOCATION = \"us-central1\"\n",
    "\n",
    "create_bucket_and_copy_data(SOURCE_BUCKET_PATH, DATASET_PATH, DATASET_PATH_LOCATION)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e62bdabd",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "### Create a BigQuery dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7b33b91",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "create_dataset_sql =f\"CREATE SCHEMA IF NOT EXISTS `{PROJECT_ID}.{DATASET_ID}` OPTIONS ( location = '{REGION}' );\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0049b41e",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "%%bigquery\n",
    "$create_dataset_sql"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "979a179f",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "### Create a BigQuery object table from the GCS location"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6dac911",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "create_object_table_sql = f\"\"\"\n",
    "      CREATE OR REPLACE EXTERNAL TABLE\n",
    "        `{PROJECT_ID}.{DATASET_ID}.{OBJECT_TABLE_NAME}`\n",
    "      WITH\n",
    "        CONNECTION DEFAULT\n",
    "        OPTIONS\n",
    "          (object_metadata = 'SIMPLE', uris = ['{DATASET_PATH}/*']);\n",
    "      \"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3572025",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "%%bigquery\n",
    "$create_object_table_sql"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45afed40",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "### Read object tables using Bigframes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c456d59b-d7cd-47db-91be-d7d153717f31",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "df = bpd.read_gbq(f\"{PROJECT_ID}.{DATASET_ID}.{OBJECT_TABLE_NAME}\", use_cache=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a46bfbb-8fc9-4e52-a804-72d4c1c31cfa",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "df"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46a0f94c-f795-401f-a784-b9c522030e2b",
   "metadata": {},
   "source": [
    "### Fetch the ads product taxonomy from the web (Interactive Advertising Bureau)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7409d1ab-ede8-4958-adc3-4df743d54f9e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "TAXONOMY_URL = 'https://raw.githubusercontent.com/InteractiveAdvertisingBureau/Taxonomies/main/Ad%20Product%20Taxonomies/Ad%20Product%20Taxonomy%202.0.tsv'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89232878-c7dc-489b-81bc-624a5fafc85f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import io\n",
    "import requests\n",
    "\n",
    "response = requests.get(TAXONOMY_URL)\n",
    "response.raise_for_status()\n",
    "ads_product_taxonomy = pd.read_csv(io.StringIO(response.text), sep='\\t', header=0)\n",
    "ads_product_taxonomy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22cc035b",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "ads_product_taxonomy_lowest_rank_list = ads_product_taxonomy['Name'].to_list()\n",
    "ads_product_taxonomy_lowest_rank_list[:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8dd063ab",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "ads_product_taxonomy_lowest_rank = '\\n'.join(ads_product_taxonomy_lowest_rank_list)\n",
    "print(ads_product_taxonomy_lowest_rank[:100])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5426730c-cdc0-4a49-9683-ce0f558fb650",
   "metadata": {},
   "source": [
    "### Create the prompt to analyze the ads banner image and generate the interpretations"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e63b538-d282-4f40-81e9-bd85ca991780",
   "metadata": {},
   "source": [
    "#### Define the BigQuery [remote function](https://cloud.google.com/python/docs/reference/bigframes/0.19.2/bigframes.session.Session#bigframes_session_Session_remote_function) to call the Gemini API"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "045941ad-bc72-47b8-837b-f2ce7c58aa22",
   "metadata": {},
   "source": [
    "By doing so, a Cloud Function will be deployed in your GCP project"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa5adf60",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Capture the output of gcloud command directly into a Python variable\n",
    "project_number = !gcloud projects list --filter=\"project_id={PROJECT_ID}\" --format=\"value(project_number)\"\n",
    "project_number = project_number[0].strip()\n",
    "default_sa = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
    "\n",
    "# Grant the role using gcloud\n",
    "!gcloud projects add-iam-policy-binding {PROJECT_ID} --member=serviceAccount:{default_sa} --role='roles/aiplatform.user' --condition=None --no-user-output-enabled"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4db45890-4124-4704-9814-c6ab3e3138e6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "@bpd.remote_function(\n",
    "    [str, str],\n",
    "    str,\n",
    "    reuse=True,\n",
    "    cloud_function_service_account=\"default\",\n",
    "    packages=[\"google-cloud-aiplatform\", \"pandas\"]\n",
    ")\n",
    "def generate_predictions(uri: str, mime_type: str) -> str:\n",
    "    from vertexai.generative_models import GenerativeModel, GenerationConfig, Part, Content, HarmCategory, HarmBlockThreshold\n",
    "    import pandas as pd\n",
    "    \n",
    "    # 1. Fetch taxonomy using Pandas\n",
    "    url = 'https://raw.githubusercontent.com/InteractiveAdvertisingBureau/Taxonomies/main/Ad%20Product%20Taxonomies/Ad%20Product%20Taxonomy%202.0.tsv'\n",
    "    try:\n",
    "        ads_product_taxonomy = pd.read_csv(url, sep='\\t', header=0)\n",
    "        ads_product_taxonomy_lowest_rank_list = ads_product_taxonomy['Name'].to_list()\n",
    "        ads_product_taxonomy_lowest_rank = '\\n'.join(ads_product_taxonomy_lowest_rank_list)\n",
    "    except Exception as e:\n",
    "        return f\"Error fetching taxonomy: {str(e)}\"\n",
    "        \n",
    "    # 2. Define prompt and schema\n",
    "    system_instructions = [\n",
    "        \"\"\"You are a marketing and advertising expert, and have in-depth knowledge of web advertising campaigns.\"\"\",\n",
    "        \"\"\"Its task is to analyze a banner in image or video format, and return various information about it.\"\"\",\n",
    "        \"\"\"You will respond in JSON format to the following fields: product, interpretation, intended_audience, classification and score.\"\"\"\n",
    "    ]\n",
    "    \n",
    "    response_schema = {\n",
    "        \"type\": \"OBJECT\",\n",
    "        \"properties\": {\n",
    "                        \"product\": {\"type\": \"STRING\"},\n",
    "                        \"interpretation\": {\"type\": \"STRING\"},\n",
    "                        \"intended_audience\": {\"type\": \"STRING\"},\n",
    "                        \"classification\": {\"type\": \"STRING\"},\n",
    "                        \"score\": {\"type\": \"INTEGER\"}\n",
    "                    },\n",
    "                    \"required\": [\"product\", \"interpretation\", \"intended_audience\", \"classification\", \"score\"],\n",
    "    }\n",
    "    \n",
    "    # 3. Create prompt\n",
    "    final_prompt = f\"\"\"\n",
    "I will give you instructions on how to obtain each piece of information.\n",
    "\n",
    "<h1>product</h1>\n",
    "What is the brand and product being promoted on the banner? Be brief in your answer, just say the name of the brand and product separated by | and nothing else.\n",
    "For example:\n",
    "Sony|BRAVIA X90K\n",
    "\n",
    "<h1>interpretation</h1>\n",
    "Generate an interpretation of the main message that the banner conveys, the product being promoted and the target audience.\n",
    "For example, for a banner showing an offer from Audible, the result would be:\n",
    "The main message of the banner is that Audible is offering a 66% discount for the first 3 months of subscription. The product being promoted is Audible, an audiobook and podcast streaming service. The target audience is anyone interested in listening to audiobooks or podcasts.\n",
    "DO NOT generate multi-line responses. Be concise, 2 to 3 sentences maximum.\n",
    "\n",
    "<h1>intended_audience</h1>\n",
    "What is the target audience for the ad? Focus on the target persona of the product being promoted.\n",
    "For example, for a banner showing an offer from Audible, the result would be:\n",
    "The target audience is anyone interested in listening to audiobooks or podcasts.\n",
    "DO NOT generate multi-line responses. Be concise, with 1 sentence.\n",
    "\n",
    "<h1>classification</h1>\n",
    "I need you to classify the banner based on a taxonomy. For classification, use the exact term, even if it is in English.\n",
    "Here is the taxonomy you should consider to classify the banner:\n",
    "<h2>taxonomy</h2>\n",
    "{ads_product_taxonomy_lowest_rank}\n",
    "</>\n",
    "For example, for a banner selling an office chair, the classification would be:\n",
    "Office Equipment and Supplies\n",
    "For example, for a banner selling a dry cleaning promotion, the classification would be:\n",
    "Laundry and Dry Cleaning Services\n",
    "DO NOT generate multi-line responses. You must ONLY generate ONE EXISTING classification IN THE TAXONOMY LIST and NO MORE WORDS!\n",
    "\n",
    "<h1>score</h1>\n",
    "Evaluate the confidence in you have in your analysis from 0 to 10.  \n",
    "If you are not sure what the banner means, you assign a lower score.  \n",
    "If it is very clear, you assign a higher score.  \n",
    "Consider this reasoning:\n",
    "Are the classifications referring to an actual quote from the content?\n",
    "Are the classifications correct, accurate and factual?\n",
    "\n",
    "<h1>Response in JSON format</h1>\n",
    "\"\"\"\n",
    "\n",
    "    model = GenerativeModel(model_name=\"gemini-2.5-flash\", system_instruction=system_instructions)\n",
    "    \n",
    "    prompt_content = Content(\n",
    "        role=\"user\",\n",
    "        parts=[\n",
    "            Part.from_uri(uri, mime_type),\n",
    "            Part.from_text(final_prompt)\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    response = model.generate_content(\n",
    "        prompt_content,\n",
    "        generation_config=GenerationConfig(\n",
    "            max_output_tokens= 8192, temperature=0.5, response_mime_type=\"application/json\", response_schema=response_schema\n",
    "        ),\n",
    "        safety_settings={\n",
    "                HarmCategory.HARM_CATEGORY_UNSPECIFIED: HarmBlockThreshold.BLOCK_ONLY_HIGH\n",
    "        }\n",
    "    )\n",
    "    \n",
    "    return response.text"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8544aa46-a853-4fb2-991b-9e2df44555ab",
   "metadata": {},
   "source": [
    "### Run predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90dc45c0-2317-406e-b148-3ecb3247bab9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "input_remote_function = df[[\"uri\"]].assign(mime_type=df[[\"content_type\"]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6b3d4d9-5c08-4b99-8729-df06a5f380d2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "result_df = df.assign(pred=input_remote_function.apply(generate_predictions, axis=1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b0281b4-747e-4ac5-b620-39fa084d3fc0",
   "metadata": {},
   "outputs": [],
   "source": [
    "result_df"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "81c389cc-35c9-4d59-a719-34261c038ef0",
   "metadata": {},
   "source": [
    "### Extract attributes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a229b7f",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "explode_columns = [\"product\", \"interpretation\", \"intended_audience\", \"classification\", \"score\"]\n",
    "\n",
    "for col in explode_columns:\n",
    "    result_df[col] = bbq.json_extract(result_df[\"pred\"], json_path=f\"$.{col}\")\n",
    "\n",
    "result_df"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96a72f2e",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "### Write the final results to BigQuery"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29d147de",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "result_df.to_gbq(destination_table=f\"{PROJECT_ID}.{DATASET_ID}.{OUTPUT_TABLE_NAME}\", if_exists=\"replace\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5a39c6c-1e72-494a-ad16-7d1af5318d87",
   "metadata": {},
   "source": [
    "### Iterate on result df to display images and prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02161610-e04f-42cf-a800-41515efd90c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def display_image_predictions(df, bucket_loc):\n",
    "    \n",
    "  gcs_client = storage.Client()\n",
    "  bucket_name = bucket_loc.split('/')[2]\n",
    "\n",
    "  for _, row in df.iterrows():\n",
    "    image_path = row['uri'].split('/', 3)[-1]\n",
    "    \n",
    "    try:\n",
    "      bucket = gcs_client.bucket(bucket_name)\n",
    "      blob = bucket.blob(image_path)\n",
    "      image_bytes = blob.download_as_bytes()\n",
    "\n",
    "      image = Image.open(io.BytesIO(image_bytes))\n",
    "\n",
    "      print(f\"Prediction for image: {row['uri']}\")\n",
    "      display(image)\n",
    "\n",
    "      prediction_data = [\n",
    "          ['Interpretation', row[\"interpretation\"]],\n",
    "          ['Intended Audience', row[\"intended_audience\"]],\n",
    "          ['Classification', row[\"classification\"]],\n",
    "          ['Score', row[\"score\"]]\n",
    "      ]\n",
    "\n",
    "      print(tabulate(prediction_data, headers=[\"Field\", \"Value\"]))\n",
    "\n",
    "    except Exception as e:\n",
    "      print(f\"Error fetching image: {e}\")\n",
    "\n",
    "display_image_predictions(result_df.to_pandas(), DATASET_PATH)"
   ]
  }
 ],
 "metadata": {
  "environment": {
   "kernel": "conda-root-py",
   "name": "workbench-notebooks.m119",
   "type": "gcloud",
   "uri": "us-docker.pkg.dev/deeplearning-platform-release/gcr.io/workbench-notebooks:m119"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel) (Local)",
   "language": "python",
   "name": "python3"
  },
  "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}