notebooks/community/model_garden/model_garden_mediapipe_face_stylizer.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ur8xi4C7S06n"
      },
      "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": "TirJ-SGQseby"
      },
      "source": [
        "# Vertex AI Model Garden MediaPipe with Face Stylizer\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td>\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_mediapipe_face_stylizer.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",
        "\n",
        "  <td>\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_mediapipe_face_stylizer.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/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_mediapipe_face_stylizer.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": "dwGLvtIeECLK"
      },
      "source": [
        "**_NOTE_**: This notebook has been tested in the following environment:\n",
        "\n",
        "* Python version = 3.9\n",
        "\n",
        "**_NOTE_**: The checkpoint and the dataset linked in this Colab are not owned or distributed by Google, and are made available by third parties. Please review the terms and conditions made available by the third parties before using the checkpoint and data."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook demonstrates how to use [MediaPipe Model Maker](https://developers.google.com/mediapipe/solutions/model_maker) to customize an on-device face stylizer model in Vertex AI Model Garden.\n",
        "\n",
        "The MediaPipe face stylizer solution provides several models you can use immediately to transform the face to the styles including (cartoon, oil painting, etc.) in your application. However, if you need to transfer the face to an unseen style not covered by the provided models, you can customize the pretrained model with your own data and MediaPipe Model Maker. This model modification tool fine-tune a portion of the model using data you provide. This method is faster than training a new model from scatch and can produce a model adapt to your specific application.\n",
        "\n",
        "The following sections show you how to use Model Maker to retrain a pre-built model for face stylization with your own data on Vertex AI, which you can then use with the MediaPipe Face Stylizer.\n",
        "\n",
        "### Objective\n",
        "\n",
        "* Customize a Face Stylizer model\n",
        "  * Convert input data to training formats\n",
        "  * Create [custom jobs](https://cloud.google.com/vertex-ai/docs/training/create-custom-job) to customize new models\n",
        "  * Export customized models\n",
        "\n",
        "* Cleanup resources\n",
        "\n",
        "### Costs\n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "* Vertex AI\n",
        "* Cloud Storage\n",
        "\n",
        "Learn about [Vertex AI\n",
        "pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage\n",
        "pricing](https://cloud.google.com/storage/pricing), and use the [Pricing\n",
        "Calculator](https://cloud.google.com/products/calculator/)\n",
        "to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KEukV6uRk_S3"
      },
      "source": [
        "## Before you begin"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "z__i0w0lCAsW"
      },
      "source": [
        "### Colab only\n",
        "Run the following commands to install dependencies and to authenticate with Google Cloud if running on Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Jvqs-ehKlaYh"
      },
      "outputs": [],
      "source": [
        "! pip3 install --upgrade pip\n",
        "\n",
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    ! pip3 install --upgrade google-cloud-aiplatform\n",
        "\n",
        "    # Automatically restart kernel after installs\n",
        "    import IPython\n",
        "\n",
        "    app = IPython.Application.instance()\n",
        "    app.kernel.do_shutdown(True)\n",
        "\n",
        "    from google.colab import auth as google_auth\n",
        "\n",
        "    google_auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WReHDGG5g0XY"
      },
      "source": [
        "#### Set your project ID\n",
        "\n",
        "**If you don't know your project ID**, see the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oM1iC_MfAts1"
      },
      "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": "region"
      },
      "source": [
        "#### Region\n",
        "\n",
        "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tTy1gX11kCJY"
      },
      "outputs": [],
      "source": [
        "REGION = \"us-central1\"  # @param {type: \"string\"}\n",
        "REGION_PREFIX = REGION.split(\"-\")[0]\n",
        "assert REGION_PREFIX in (\n",
        "    \"us\",\n",
        "    \"europe\",\n",
        "    \"asia\",\n",
        "), f'{REGION} is not supported. It must be prefixed by \"us\", \"asia\", or \"europe\".'"
      ]
    },
    {
      "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": [
        "**Only 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 {REGION} -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 json\n",
        "import os\n",
        "from datetime import datetime\n",
        "\n",
        "from google.cloud import aiplatform"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "init_aip:mbsdk,all"
      },
      "source": [
        "### Initialize Vertex AI SDK for Python\n",
        "\n",
        "Initialize the Vertex AI SDK for Python for your project."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9wExiMUxFk91"
      },
      "outputs": [],
      "source": [
        "now = datetime.now().strftime(\"%Y%m%d-%H%M%S\")\n",
        "\n",
        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temp/%s\" % now)\n",
        "\n",
        "\n",
        "EXPORTED_MODEL_OUTPUT_DIRECTORY = os.path.join(STAGING_BUCKET, \"model\")\n",
        "EXPORTED_MODEL_OUTPUT_FILE = os.path.join(\n",
        "    EXPORTED_MODEL_OUTPUT_DIRECTORY, \"model.tflite\"\n",
        ")\n",
        "\n",
        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "n6IFz75WGCam"
      },
      "source": [
        "### Define training machine specs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "riG_qUokg0XZ"
      },
      "outputs": [],
      "source": [
        "TRAINING_JOB_DISPLAY_NAME = \"mediapipe_face_stylizer_%s\" % now\n",
        "TRAINING_CONTAINER = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/mediapipe-train\"\n",
        "TRAINING_MACHINE_TYPE = \"n1-highmem-16\"\n",
        "TRAINING_ACCELERATOR_TYPE = \"NVIDIA_TESLA_V100\"\n",
        "TRAINING_ACCELERATOR_COUNT = 2"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-rsdAcBV-vlf"
      },
      "source": [
        "## Train your customized models"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LmZ4efOd-sak"
      },
      "source": [
        "### Prepare input data for training\n",
        "\n",
        "Retraining the face stylizer model requires user to provide a single stylized face image. The stylized face is expected to be forward facing with visible left right eyes and mouth. The face should only have minor rotation, i.e. less than 30 degress around the yaw, pitch, and roll axes.\n",
        "\n",
        "You can upload an image to Google Cloud Storage or use our [provided example](https://storage.googleapis.com/mediapipe-assets/face_stylizer_style_color_sketch.jpg)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IndQ_m6ddUEM"
      },
      "outputs": [],
      "source": [
        "training_data_path = \"gs://mediapipe-assets/face_stylizer_style_color_sketch.jpg\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aaff6f5be7f6"
      },
      "source": [
        "# Retrain model\n",
        "\n",
        "Once you have provided an input image, you can begin retraining the face stylizer model to adapt to the new style. This type of model modification is called transfer learning. The instructions below use the data prepared in the previous section to retrain a face stylizer model to apply cartoon style to the raw human face.\n",
        "\n",
        "**_NOTE_**: For this type of model, the retraining process causes the model to forget any style it can apply before. Once the retraining is complete, the new model can only apply the new style defined by the new stylized image.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kySYKb3-hnr6"
      },
      "source": [
        "## Set retraining options\n",
        "There are a few required settings to run a retraining aside from your training dataset:\n",
        "\n",
        "* **Swap layers:** The `swap_layers` parameter is used to determine how to mix the latent code layers between the learned style and the raw face images. The latent code is represented as a tensor of shape [1, 12, 512]. The second dimension of the latent code tensor is called the layer. The face stylizer mixes the learned style and raw face images by generating a weighted sum of the two latent codes on the swap layers. The swap layers are therefore integers within [1, 12]. The more layers are set, the more style will be applied to the output image. Although there is no explicit mapping between the style semantics and the layer index, the shallow layers, e.g. 8, 9, represent the global features of the face, while the deep layers, e.g. 10, 11, represent the fine-grained features of the face. The output stylized image is sensitive to the setting of swap layers. By default, it is set to [8, 9, 10, 11].\n",
        "* **Learning rate and epochs:** Use `learning_rate` and epochs` to specify the these two hyperparameters. learning_rate is set to 4e-4 by default. epochs defines the number of iterations to fine-tune the BlazeStyleGAN model and are set to 100 by default. The lower the learning rate is, the greater the epochs is expected to retrain the model to converge.\n",
        "* **Batch size:** The `batch_size` is used to define the number of latent code samples we sample around the latent code extracted by the encoder with the input image. The batch of latent codes are used to fine-tune the decoder. The greater the batch size usually yield to better performance. It is also limited by the hardware memory. For A100 GPU, the maximum batch size is 8. For P100 and T4 GPU, the maximum batch size is 2.\n",
        "\n",
        "Further more advanced parameters that you can configure are `alpha`, `perception_loss_weight`, `adv_loss_weight`, `beta_1` and `beta_2`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "um_XKbmpTaHx"
      },
      "outputs": [],
      "source": [
        "# The layers of feature to be interpolated between encoding features and\n",
        "# StyleGAN input features.\n",
        "swap_layers: str = \"[8, 9, 10, 11]\"  # @param {type:\"string\"}\n",
        "# The learning rate to use for gradient descent training.\n",
        "learning_rate: float = 0.0001  # @param {type:\"number\"}\n",
        "# Number of training iterations over the dataset.\n",
        "epochs: int = 100  # @param {type:\"slider\", min:0, max:100, step:1}\n",
        "# Batch size for training.\n",
        "batch_size: int = 2  # @param {type:\"number\"}\n",
        "\n",
        "\n",
        "# Other supported options\n",
        "\n",
        "# Weighting coefficient of style latent for swapping layer interpolation.\n",
        "# Its valid range is [0, 1]. The greater weight means stronger style is\n",
        "# applied to the output image. Expect to set it to a small value,\n",
        "# i.e. < 0.1.\n",
        "alpha: float = 0.1  # @param {type:\"number\"}\n",
        "\n",
        "# Weighting coefficients of image perception quality loss. It contains three\n",
        "# coefficients, l1, content, and style which control the difference between the\n",
        "# generated image and raw input image, the content difference between generated\n",
        "# face and raw input face, and the how similar the style between the generated\n",
        "# image and raw input image. Users can increase the style weight to enforce\n",
        "# stronger style or the content weight to reserve more raw input face details.\n",
        "# Weight for L1 loss.\n",
        "perception_loss_l1: float = 0.5  # @param {type:\"number\"}\n",
        "# Weight for content loss.\n",
        "perception_loss_content: float = 4.0  # @param {type:\"number\"}\n",
        "# Weight for stlye loss.\n",
        "perception_loss_style: float = 1.0  # @param {type:\"number\"}\n",
        "\n",
        "# Weighting coeffcieint of adversarial loss versus image perceptual quality loss.\n",
        "# This hyperparameter is used to control the realism of the generated image. It\n",
        "# expects a small value, i.e. < 0.2.\n",
        "adv_loss_weight: float = 0.2  # @param {type:\"number\"}\n",
        "# beta_1 used in tf.keras.optimizers.Adam.\n",
        "beta_1: float = 0.0  # @param {type:\"number\"}\n",
        "# beta_2 used in tf.keras.optimizers.Adam.\n",
        "beta_2: float = 0.99  # @param {type:\"number\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HwcCjwlBTQIz"
      },
      "source": [
        "### Run retraining\n",
        "With your training dataset and retraining options prepared, you are ready to start the retraining process. This process requires running on GPU and can take a few minutes to a few hours depending on your available compute resources. On Vertex AI with GPU processing, the example retraining below takes about 2 minutes.\n",
        "\n",
        "To begin the fine-tuning process, use the following code:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aec22792ee84"
      },
      "outputs": [],
      "source": [
        "model_export_path = EXPORTED_MODEL_OUTPUT_DIRECTORY\n",
        "\n",
        "worker_pool_specs = [\n",
        "    {\n",
        "        \"machine_spec\": {\n",
        "            \"machine_type\": TRAINING_MACHINE_TYPE,\n",
        "            \"accelerator_type\": TRAINING_ACCELERATOR_TYPE,\n",
        "            \"accelerator_count\": TRAINING_ACCELERATOR_COUNT,\n",
        "        },\n",
        "        \"replica_count\": 1,\n",
        "        \"container_spec\": {\n",
        "            \"image_uri\": TRAINING_CONTAINER,\n",
        "            \"command\": [],\n",
        "            \"args\": [\n",
        "                \"--task_name=face_stylizer\",\n",
        "                \"--training_data_path=%s\" % training_data_path,\n",
        "                \"--model_export_path=%s\" % model_export_path,\n",
        "                \"--evaluation_result_path=%s\" % model_export_path,\n",
        "                \"--hparams=%s\"\n",
        "                % json.dumps(\n",
        "                    {\n",
        "                        \"learning_rate\": learning_rate,\n",
        "                        \"batch_size\": batch_size,\n",
        "                        \"epochs\": epochs,\n",
        "                        \"beta_1\": beta_1,\n",
        "                        \"beta_2\": beta_2,\n",
        "                    }\n",
        "                ),\n",
        "                \"--model_options=%s\"\n",
        "                % json.dumps(\n",
        "                    {\n",
        "                        \"swap_layers\": json.loads(swap_layers),\n",
        "                        \"alpha\": alpha,\n",
        "                        \"perception_loss_l1\": perception_loss_l1,\n",
        "                        \"perception_loss_content\": perception_loss_content,\n",
        "                        \"perception_loss_style\": perception_loss_style,\n",
        "                        \"adv_loss_weight\": adv_loss_weight,\n",
        "                    }\n",
        "                ),\n",
        "            ],\n",
        "        },\n",
        "    }\n",
        "]\n",
        "\n",
        "training_job = aiplatform.CustomJob(\n",
        "    display_name=TRAINING_JOB_DISPLAY_NAME,\n",
        "    project=PROJECT_ID,\n",
        "    worker_pool_specs=worker_pool_specs,\n",
        "    staging_bucket=STAGING_BUCKET,\n",
        ")\n",
        "\n",
        "training_job.run()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "g0BGaofgsMsy"
      },
      "source": [
        "## Export model\n",
        "After retraining the model, you can save the Tensorflow Lite model and integrate it with your on-device application by following the [Face stylization task guide](https://developers.google.com/mediapipe/solutions/vision/face_stylizer)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NYuQowyZEtxK"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "def copy_model(model_source, model_dest):\n",
        "    ! gsutil cp {model_source} {model_dest}\n",
        "\n",
        "copy_model(EXPORTED_MODEL_OUTPUT_FILE, \"face_stylizer.task\")\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import files\n",
        "\n",
        "    files.download(\"face_stylizer.task\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kkH2nrpdp4sp"
      },
      "source": [
        "## Clean up"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ax6vQVZhp9pR"
      },
      "outputs": [],
      "source": [
        "# Delete training data and jobs.\n",
        "if training_job.list(filter=f'display_name=\"{TRAINING_JOB_DISPLAY_NAME}\"'):\n",
        "    training_job.delete()\n",
        "\n",
        "!gsutil rm -r {STAGING_BUCKET}"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "model_garden_mediapipe_face_stylizer.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}