{ "cells": [ { "cell_type": "code", "execution_count": null, "metadata": { "id": "ur8xi4C7S06n" }, "outputs": [], "source": [ "# Copyright 2024 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": [ "# Product attributes extraction and detailed descriptions from images using Gemini 2.0\n", "\n", "\n", " \n", " \n", " \n", " \n", "
\n", " \n", " \"Google
Open in Colab\n", " \n", " 		\n", " \n", " \"Google
Open in Colab Enterprise\n", " \n", " 		\n", " \n", " \"Vertex
Open in Workbench\n", " \n", " 		\n", " \n", " \"GitHub
View on GitHub\n", " \n", "
\n", "\n", "
\n", "\n", "Share to:\n", "\n", "\n", " \"LinkedIn\n", "\n", "\n", "\n", " \"Bluesky\n", "\n", "\n", "\n", " \"X\n", "\n", "\n", "\n", " \"Reddit\n", "\n", "\n", "\n", " \"Facebook\n", " " ] }, { "cell_type": "markdown", "metadata": { "id": "84f0f73a0f76" }, "source": [ "| | |\n", "|-|-|\n", "|Author(s) | [Tianli Yu](https://github.com/tianli) |" ] }, { "cell_type": "markdown", "metadata": { "id": "tvgnzT1CKxrO" }, "source": [ "## Overview\n", "\n", "This Colab teaches how to build a general agent (on top of Gemini) to extract different product attributes or detailed product descriptions from an image input. It also introduces a prompting technique called \"self-correcting prompt\" where you can ask the model to check and verify the result by itself (all in one single prompt). Self-correcting prompt proves to improve the overall quality of the agent's output.\n", "\n", "In the following sections we will:\n", "\n", "* Write the necessary image loading and parsing library.\n", "* Create a product image agent.\n", "* Run the agent on a set of examples for tasks like product image description and attribute extraction." ] }, { "cell_type": "markdown", "metadata": { "id": "61RBz8LLbxCR" }, "source": [ "## Get started" ] }, { "cell_type": "markdown", "metadata": { "id": "No17Cw5hgx12" }, "source": [ "### Install Vertex AI SDK and other required packages" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "tFy3H3aPgx12" }, "outputs": [], "source": [ "%pip install --upgrade --user --quiet google-cloud-aiplatform" ] }, { "cell_type": "markdown", "metadata": { "id": "R5Xep4W9lq-Z" }, "source": [ "### Restart runtime\n", "\n", "To use the newly installed packages in this Jupyter runtime, you must restart the runtime. You can do this by running the cell below, which restarts the current kernel.\n", "\n", "The restart might take a minute or longer. After it's restarted, continue to the next step." ] }, { "cell_type": "code", "execution_count": 1, "metadata": { "id": "XRvKdaPDTznN" }, "outputs": [ { "data": { "text/plain": [ "{'status': 'ok', 'restart': True}" ] }, "execution_count": 1, "metadata": {}, "output_type": "execute_result" } ], "source": [ "import IPython\n", "\n", "app = IPython.Application.instance()\n", "app.kernel.do_shutdown(True)" ] }, { "cell_type": "markdown", "metadata": { "id": "SbmM4z7FOBpM" }, "source": [ "
\n", "⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️\n", "
" ] }, { "cell_type": "markdown", "metadata": { "id": "dmWOrTJ3gx13" }, "source": [ "### Authenticate your notebook environment (Colab only)\n", "\n", "If you're running this notebook on Google Colab, run the cell below to authenticate your environment." ] }, { "cell_type": "code", "execution_count": 1, "metadata": { "id": "NyKGtVQjgx13" }, "outputs": [], "source": [ "import sys\n", "\n", "if \"google.colab\" in sys.modules:\n", " from google.colab import auth\n", "\n", " auth.authenticate_user()" ] }, { "cell_type": "markdown", "metadata": { "id": "DF4l8DTdWgPY" }, "source": [ "### Set Google Cloud project information and initialize Vertex AI SDK\n", "\n", "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\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": 2, "metadata": { "id": "Nqwi-5ufWp_B" }, "outputs": [], "source": [ "PROJECT_ID = \"YOUR_PROJECT_ID\" # @param {type:\"string\"}\n", "LOCATION = \"us-central1\" # @param {type:\"string\"}\n", "\n", "\n", "import vertexai\n", "\n", "vertexai.init(project=PROJECT_ID, location=LOCATION)" ] }, { "cell_type": "markdown", "metadata": { "id": "EdvJRUWRNGHE" }, "source": [ "## Build a Product Image Agent" ] }, { "cell_type": "code", "execution_count": 3, "metadata": { "cellView": "form", "id": "xJQczGZX3FfW" }, "outputs": [], "source": [ "# @title Image loading and parsing library\n", "import json\n", "\n", "import requests\n", "from vertexai.preview.generative_models import Part\n", "\n", "\n", "def get_mime_from_uri(image_uri: str) -> str:\n", " \"\"\"Get the mime type from the image uri.\"\"\"\n", " if image_uri.endswith(\".png\"):\n", " return \"image/png\"\n", " elif image_uri.endswith(\".gif\"):\n", " return \"image/gif\"\n", " else:\n", " # Assume JPEG as the default mime\n", " return \"image/jpeg\"\n", "\n", "\n", "def load_image_part_from_uri(image_uri: str) -> None | Part:\n", " \"\"\"Load image to a prompt Part from a remote or local file URI.\n", "\n", " Args:\n", " image_uri (str): The uri or the local file path for the image.\n", "\n", " Returns:\n", " The prompt Part object.\n", " \"\"\"\n", " if image_uri.startswith(\"gs://\"):\n", " return Part.from_uri(image_uri, mime_type=get_mime_from_uri(image_uri))\n", " elif image_uri.startswith(\"http://\") or image_uri.startswith(\"https://\"):\n", " response = requests.get(image_uri)\n", " if response.status_code == 200:\n", " image_bytes = response.content\n", " return Part.from_data(image_bytes, mime_type=get_mime_from_uri(image_uri))\n", " else:\n", " print(\n", " f\"Fetch image failed for {image_uri}, status code: {response.status_code}\"\n", " )\n", " return None\n", " else:\n", " image_bytes = open(image_uri, \"rb\").read()\n", " return Part.from_data(image_bytes, mime_type=get_mime_from_uri(image_uri))\n", "\n", "\n", "def parse_json_from_markdown(answer: str) -> str:\n", " \"\"\"Parse the json from the markdown answer.\n", "\n", " Args:\n", " answer (str): The markdown answer from the model.\n", "\n", " Returns:\n", " A parsed json dictionary}\n", " \"\"\"\n", " lines = answer.split(\"```\")\n", " try:\n", " # Tries to parse the last json in the answer.\n", " answer = lines[-2]\n", " if answer.startswith(\"json\"):\n", " answer = answer[4:]\n", " result = json.loads(answer)\n", " except json.JSONDecodeError:\n", " # Falls back to the first json in the answer.\n", " answer = lines[1]\n", " if answer.startswith(\"json\"):\n", " answer = answer[4:]\n", " result = json.loads(answer)\n", " return json.dumps(result)" ] }, { "cell_type": "code", "execution_count": 44, "metadata": { "id": "3nThlwYK_tN2" }, "outputs": [], "source": [ "# @title The ProductImageAgent Class.\n", "\n", "\n", "from IPython.display import display\n", "import ipywidgets as widgets\n", "from vertexai.preview.generative_models import GenerationConfig, GenerativeModel, Part\n", "\n", "\n", "class ProductImageAgent:\n", " \"\"\"An agent that wraps around Gemini 2.0 to extract product attributes from\n", " images.\n", "\n", " Args:\n", " gemini_model_version (str): The version string of the Gemini 2.0 model.\n", " gemini-2.0-flash or gemini-2.0-flash\n", " temperature (float): The temperature of the model. Defaults to 1.0.\n", " max_output_tokens (int): The maximum number of output tokens. Defaults to\n", " 8192.\n", " \"\"\"\n", "\n", " def __init__(\n", " self,\n", " gemini_model_version: str = \"gemini-2.0-flash\",\n", " temperature: float = 0.0,\n", " max_output_tokens: int = 8192,\n", " ):\n", " config = GenerationConfig(\n", " temperature=temperature, max_output_tokens=max_output_tokens\n", " )\n", "\n", " # System instructions, add any common instructions here.\n", " sys_inst = \"\"\"\n", " As an assistant for an online retailer, your task is to recognize\n", " attributes from the provided product image.\n", " If an attribute vocabulary is provided, please only select attribute values\n", " in that vocabulary. You answer should be strictly consistent with what's in\n", " the image. If any attributes do not exist in the image, please\n", " return null for that attribute.\n", " \"\"\"\n", " self.gemini_model = GenerativeModel(\n", " gemini_model_version, generation_config=config, system_instruction=sys_inst\n", " )\n", "\n", " def get_detailed_description(self, image_uri: str, debug: bool = False) -> str:\n", " \"\"\"Generates the detailed product description from an image.\n", "\n", " Args:\n", " image_uri: The url to the image, can be a local file path, or a url\n", " from the web or gcs.\n", "\n", " Returns:\n", " The generated detailed description from the model response.\n", " \"\"\"\n", " image_part = load_image_part_from_uri(image_uri)\n", " if image_part is None:\n", " return \"\"\n", " prompt = \"\"\"\n", " Please write a complete and detailed product description for the\n", " above product image. The length of the description should be at least\n", " 200 words.\n", " \"\"\"\n", " if debug:\n", " print(\"====== Begin Debug Info ======\")\n", " preview = widgets.HTML(value=f'')\n", " display(preview)\n", " print(f\"Prompt:\\n{prompt}\")\n", " print(\"====== End Debug Info ======\")\n", "\n", " model_response = self.gemini_model.generate_content([image_part, prompt])\n", " return model_response.text\n", "\n", " def get_attributes(\n", " self,\n", " image_uri: str,\n", " product_category: None | str = None,\n", " vocabulary_json: None | str = None,\n", " debug: bool = False,\n", " ) -> str:\n", " \"\"\"Generates the product attributes from an image.\n", "\n", " Args:\n", " image_uri (str): The uri of the product image to generate attributes.\n", " vocabulary_json (str): A json string list all the attribute names and\n", " their possible vocabulary.\n", "\n", " Returns:\n", " The product attribute json string from the model response.\n", " \"\"\"\n", " image_part = load_image_part_from_uri(image_uri)\n", " if image_part is None:\n", " return \"\"\n", " if product_category:\n", " prompt = f\"\"\"\n", " The above image is a product image from the {product_category}\n", " category.\n", " Please list all the relevant attributes in the {product_category}\n", " category for the main product in the above image and return a list of\n", " key-value pairs in json format.\n", " \"\"\"\n", " else:\n", " prompt = \"\"\"\n", " Please recognize the main product's all relevant attributes in the\n", " above image and return a list of key-value pairs in json format.\n", " \"\"\"\n", " if vocabulary_json:\n", " prompt += f\"\"\"\n", " Please using only the vocabulary defined in the following json:\n", " {vocabulary_json}\n", " For each key, you should select the most appropriate attribute value\n", " from its corresponding vocabulary list and returns one value\n", " for each attribute key.\n", " You can return null for that key if no attributes match.\n", " \"\"\"\n", " if debug:\n", " print(\"====== Begin Debug Info ======\")\n", " preview = widgets.HTML(value=f'')\n", " display(preview)\n", " print(f\"Prompt:\\n{prompt}\")\n", " print(\"====== End Debug Info ======\")\n", "\n", " model_response = self.gemini_model.generate_content([image_part, prompt])\n", " return model_response.text\n", "\n", " def get_attributes_self_correcting_prompt(\n", " self,\n", " image_uri: str,\n", " product_category: None | str = None,\n", " vocabulary_json: None | str = None,\n", " debug: bool = False,\n", " ) -> str:\n", " \"\"\"Generates the product attributes from an image using self-correcting prompt.\n", "\n", " Args:\n", " image_uri (str): The uri of the product image to generate attributes.\n", " vocabulary_json (str): A json string list all the attribute names and\n", " their possible vocabulary.\n", "\n", " Returns:\n", " The product attribute json string from the parsed model response.\n", " \"\"\"\n", " image_part = load_image_part_from_uri(image_uri)\n", " if image_part is None:\n", " return \"\"\n", " if product_category:\n", " prompt = f\"\"\"\n", " The above image is a product image from the {product_category}\n", " category.\n", " First please list all the relevant attributes in the\n", " {product_category} category for the main product in the above\n", " image and return a list of key-value pairs in json format.\n", " \"\"\"\n", " else:\n", " prompt = \"\"\"\n", " First Please recognize the main product's all relevant attributes in\n", " the above image and return a list of key-value pairs in json\n", " format.\n", " \"\"\"\n", "\n", " if vocabulary_json:\n", " prompt += f\"\"\"\n", " Please using only the vocabulary defined in the following json:\n", " {vocabulary_json}\n", " For each key, you should select the most appropriate attribute value\n", " from its corresponding vocabulary list and return the attribute key-\n", " value pair. You can return null for that key if no attributes match.\n", " \"\"\"\n", "\n", " # Adding the self-correction instructions.\n", " prompt += \"\"\"\n", " Next, treat the returned json as the result generated by a different\n", " model, rate each key-value pair as \"correct\" or \"wrong\" based on the\n", " same image. You can output in a format like \"key - value: correct (or\n", " wrong)\".\n", " Then, based on this evaluation, please update all the attributes that\n", " are corrected in the final json output.\n", " Please use markdown to annotate different json in your output.\n", " \"\"\"\n", " model_response = self.gemini_model.generate_content([image_part, prompt])\n", "\n", " if debug:\n", " print(\"====== Begin Debug Info ======\")\n", " preview = widgets.HTML(value=f'')\n", " display(preview)\n", " print(f\"Prompt:\\n{prompt}\\n\")\n", " print(f\"Response:\\n{model_response.candidates[0].content.parts[0].text}\\n\")\n", " print(\"====== End Debug Info ======\")\n", "\n", " # Parse the model_response to get the final json.\n", " return parse_json_from_markdown(\n", " model_response.candidates[0].content.parts[0].text\n", " )\n", "\n", "\n", "# Creates the agent.\n", "product_agent = ProductImageAgent(\n", " gemini_model_version=\"gemini-2.0-flash-preview-0514\", temperature=0\n", ")" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "OCF1YAGzASQw" }, "outputs": [], "source": [ "# @title Example 1: Generate detailed product description from an image.\n", "image_uri = \"IMAGE_URI\" # @param {type:\"string\"}\n", "\n", "product_agent.get_detailed_description(image_uri, True)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "fLwaCYeHl4An" }, "outputs": [], "source": [ "# @title Example 2: Get product attribute json from an image -- open vocabulary.\n", "image_uri = \"IMAGE_URI\" # @param {type:\"string\"}\n", "\n", "# Open vocabulary.\n", "attribute_json = product_agent.get_attributes(image_uri, debug=True)\n", "print(f\"Open vocabulary attributes:\\n{attribute_json}\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "Q-afCK3lY0U3" }, "outputs": [], "source": [ "# @title Example 3: Get product attribute json from an image -- closed vocabulary.\n", "image_uri = \"IMAGE_URI\" # @param {type:\"string\"}\n", "\n", "# Closed vocabulary.\n", "vocabulary = \"\"\"\n", "{\"Pattern\": [\"Animal\", \"Letter\", \"Plaid\", \"Plain\", \"Polka Dot\", \"Quilted\", \"Striped\", \"Tie Dye\", \"Tropical\", \"Zebra\", \"Block\", \"Rainbow\", \"Floral\"], \"Toe\": [\"Almond Toe\", \"Cap Toe\", \"Closed Toe\", \"Peep Toe\", \"Point Toe\", \"Pointed Toe\", \"Round Toe\", \"Square Toe\", \"Toe Post\", \"Open Toe\"], \"Style\": [\"Ballet\", \"Bandage\", \"Basics\", \"Casual\", \"Classic\", \"Cute\", \"Elegant\", \"Formal\", \"Modern\", \"Motorcycle\", \"Retro\", \"Sexy\", \"Boho\", \"Modest\", \"Comfort\", \"Minimalist\"], \"Strap Type\": [\"Adjustable\", \"Ankle cuff\", \"Ankle straps\", \"Chain\", \"Convertible\", \"Criss Cross\", \"D'orsay\", \"Double Handle\", \"Flowers\", \"Gladiator\", \"Lace Up\", \"Mary Jane\", \"Ring\", \"Slingbacks\", \"Strappy\", \"T strap\", \"Zipper\", \"Elastic\", \"Velcro\", \"Ankle Strap\"], \"Heels\": [\"Chunky\", \"Cork\", \"Espadrilles\", \"Flat\", \"Flatform\", \"Platform\", \"Stiletto\", \"Cone Heel\", \"Kitten Heels\", \"Hidden Wedge\", \"Wedges\", \"Pyramid\"], \"Closure Type\": [\"Back Zipper\", \"Buckle\", \"Zipper\", \"Magnet\", \"Slip on\", \"Hook Loop\", \"Lace-up\", \"Flap\"]}\n", "\"\"\"\n", "\n", "attribute_json = product_agent.get_attributes(\n", " image_uri, vocabulary_json=vocabulary, debug=True\n", ")\n", "print(f\"Closed vocabulary attributes\\n{attribute_json}\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "aVCiddCOfYZ0" }, "outputs": [], "source": [ "# @title Example 4: Get product attribute json from an image using self-correcting prompt -- open vocabulary\n", "image_uri = \"IMAGE_URI\" # @param {type:\"string\"}\n", "\n", "# Open vocabulary.\n", "attribute_json = product_agent.get_attributes_self_correcting_prompt(\n", " image_uri, product_category=\"Shoes\", debug=True\n", ")\n", "print(f\"Open vocabulary attributes:\\n{attribute_json}\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "rJk--yXekwcQ" }, "outputs": [], "source": [ "# @title Example 5: Get product attribute json from an image using self-correcting prompt -- closed vocabulary\n", "image_uri = \"IMAGE_URI\" # @param {type:\"string\"}\n", "\n", "# Closed vocabulary.\n", "vocabulary = \"\"\"\n", "{\"Pattern\": [\"Animal\", \"Letter\", \"Plaid\", \"Plain\", \"Polka Dot\", \"Quilted\", \"Striped\", \"Tie Dye\", \"Tropical\", \"Zebra\", \"Block\", \"Rainbow\", \"Floral\"], \"Toe\": [\"Almond Toe\", \"Cap Toe\", \"Closed Toe\", \"Peep Toe\", \"Point Toe\", \"Pointed Toe\", \"Round Toe\", \"Square Toe\", \"Toe Post\", \"Open Toe\"], \"Style\": [\"Ballet\", \"Bandage\", \"Basics\", \"Casual\", \"Classic\", \"Cute\", \"Elegant\", \"Formal\", \"Modern\", \"Motorcycle\", \"Retro\", \"Sexy\", \"Boho\", \"Modest\", \"Comfort\", \"Minimalist\"], \"Strap Type\": [\"Adjustable\", \"Ankle cuff\", \"Ankle straps\", \"Chain\", \"Convertible\", \"Criss Cross\", \"D'orsay\", \"Double Handle\", \"Flowers\", \"Gladiator\", \"Lace Up\", \"Mary Jane\", \"Ring\", \"Slingbacks\", \"Strappy\", \"T strap\", \"Zipper\", \"Elastic\", \"Velcro\", \"Ankle Strap\"], \"Heels\": [\"Chunky\", \"Cork\", \"Espadrilles\", \"Flat\", \"Flatform\", \"Platform\", \"Stiletto\", \"Cone Heel\", \"Kitten Heels\", \"Hidden Wedge\", \"Wedges\", \"Pyramid\"], \"Closure Type\": [\"Back Zipper\", \"Buckle\", \"Zipper\", \"Magnet\", \"Slip on\", \"Hook Loop\", \"Lace-up\", \"Flap\"]}\n", "\"\"\"\n", "attribute_json = product_agent.get_attributes_self_correcting_prompt(\n", " image_uri, vocabulary_json=vocabulary, debug=True\n", ")\n", "print(f\"Closed vocabulary attributes\\n{attribute_json}\")" ] } ], "metadata": { "colab": { "name": "product_attributes_extraction.ipynb", "toc_visible": true }, "kernelspec": { "display_name": "Python 3", "name": "python3" } }, "nbformat": 4, "nbformat_minor": 0 }