Merge pull request google-ai-edge#44 from googlesamples/python_image_…

…segmentation

Added image segmentation for python example

examples/image_segmentation/python/image_segmentation.ipynb

@@ -0,0 +1,312 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "h2q27gKz1H20"
+      },
+      "source": [
+        "##### Copyright 2023 The MediaPipe Authors. All Rights Reserved."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "TUfAcER1oUS6",
+        "cellView": "form"
+      },
+      "outputs": [],
+      "source": [
+        "#@title 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": "L_cQX8dWu4Dv"
+      },
+      "source": [
+        "# Image Segmenter"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "O6PN9FvIx614"
+      },
+      "source": [
+        "## Imports and Setup\n",
+        "Let's start with the base imports for running this MediaPipe sample.\n",
+        "\n",
+        "*Notes:*\n",
+        "\n",
+        "* *If you see an error about flatbuffers incompatibility, it's fine to ignore it. MediaPipe requires a newer version of flatbuffers (v2), which is incompatible with the older version of Tensorflow (v2.9) currently preinstalled on Colab.*\n",
+        "\n",
+        "* *If you install MediaPipe outside of Colab, you only need to run pip install mediapipe. It isn't necessary to explicitly install flatbuffers*"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "gxbHBsF-8Y_l"
+      },
+      "outputs": [],
+      "source": [
+        "!pip install -q flatbuffers==2.0.0\n",
+        "!pip install -q mediapipe==0.9.0.1"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "a49D7h4TVmru"
+      },
+      "source": [
+        "## Download the image segmenter model"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "The next thing you will need to do is download the image segmentation model that will be used for this demo. In this case you will use the DeepLab v3 model."
+      ],
+      "metadata": {
+        "id": "QHXsuWDxpOj4"
+      }
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 3,
+      "metadata": {
+        "id": "OMjuVQiDYJKF"
+      },
+      "outputs": [],
+      "source": [
+        "#@title Start downloading here.\n",
+        "!wget -O deeplabv3.tflite -q https://storage.googleapis.com/mediapipe-assets/deeplabv3.tflite"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "83PEJNp9yPBU"
+      },
+      "source": [
+        "## Download a test image\n",
+        "\n",
+        "After downloading the model, it's time to grab an image that you can use for testing! It's worth noting that while this is working with a single image, you can download a collection of images to store in the `IMAGE_FILENAMES` array."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 4,
+      "metadata": {
+        "id": "tzXuqyIBlXer"
+      },
+      "outputs": [],
+      "source": [
+        "import urllib\n",
+        "\n",
+        "IMAGE_FILENAMES = ['segmentation_input_rotation0.jpg']\n",
+        "\n",
+        "for name in IMAGE_FILENAMES:\n",
+        "  url = f'https://storage.googleapis.com/mediapipe-assets/{name}'\n",
+        "  urllib.request.urlretrieve(url, name)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "P8XRmapjySMN"
+      },
+      "source": [
+        "## Preview the downloaded image"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "With the test image downloaded, go ahead and display it."
+      ],
+      "metadata": {
+        "id": "Eu_q-Z03r-ed"
+      }
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "8rjHk72-lmHX"
+      },
+      "outputs": [],
+      "source": [
+        "import cv2\n",
+        "from google.colab.patches import cv2_imshow\n",
+        "import math\n",
+        "\n",
+        "# Height and width that will be used by the model\n",
+        "DESIRED_HEIGHT = 480\n",
+        "DESIRED_WIDTH = 480\n",
+        "\n",
+        "# Performs resizing and showing the image\n",
+        "def resize_and_show(image):\n",
+        "  h, w = image.shape[:2]\n",
+        "  if h < w:\n",
+        "    img = cv2.resize(image, (DESIRED_WIDTH, math.floor(h/(w/DESIRED_WIDTH))))\n",
+        "  else:\n",
+        "    img = cv2.resize(image, (math.floor(w/(h/DESIRED_HEIGHT)), DESIRED_HEIGHT))\n",
+        "  cv2_imshow(img)\n",
+        "\n",
+        "\n",
+        "# Preview the image(s)\n",
+        "images = {name: cv2.imread(name) for name in IMAGE_FILENAMES}\n",
+        "for name, image in images.items():\n",
+        "  print(name)\n",
+        "  resize_and_show(image)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "Iy4r2_ePylIa"
+      },
+      "source": [
+        "## Running inference and visualizing the results\n",
+        "To run inference using the image segmentation MediaPipe Task, you will need to initialize the `ImageSegmenter` using the model. This example will separate the background and foreground of the image and apply separate colors for them to highlight where each distinctive area exists. Image segmentation here will use a category mask, which applies a category to each found item based on confidence."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Yl_Oiye4mUuo"
+      },
+      "outputs": [],
+      "source": [
+        "import numpy as np\n",
+        "import mediapipe as mp\n",
+        "\n",
+        "from mediapipe.python._framework_bindings import image\n",
+        "from mediapipe.python._framework_bindings import image_frame\n",
+        "from mediapipe.tasks import python\n",
+        "from mediapipe.tasks.python import vision\n",
+        "\n",
+        "\n",
+        "BG_COLOR = (192, 192, 192) # gray\n",
+        "MASK_COLOR = (255, 255, 255) # white\n",
+        "\n",
+        "OutputType = vision.ImageSegmenterOptions.OutputType\n",
+        "Activation = vision.ImageSegmenterOptions.Activation\n",
+        "\n",
+        "# Create the options that will be used for ImageSegmenter\n",
+        "base_options = python.BaseOptions(model_asset_path='deeplabv3.tflite')\n",
+        "options = vision.ImageSegmenterOptions(base_options=base_options,\n",
+        "                                              output_type=OutputType.CATEGORY_MASK)\n",
+        "\n",
+        "# Create the image segmenter\n",
+        "with vision.ImageSegmenter.create_from_options(options) as segmenter:\n",
+        "\n",
+        "  # Loop through demo image(s)\n",
+        "  for image_file_name in IMAGE_FILENAMES:\n",
+        "\n",
+        "    # Create the MediaPipe image file that will be segmented\n",
+        "    image = mp.Image.create_from_file(image_file_name)\n",
+        "\n",
+        "    # Retrieve the masks for the segmented image\n",
+        "    category_masks = segmenter.segment(image)\n",
+        "\n",
+        "    # Generate solid color images for showing the output segmentation mask.\n",
+        "    image_data = image.numpy_view()\n",
+        "    fg_image = np.zeros(image_data.shape, dtype=np.uint8)\n",
+        "    fg_image[:] = MASK_COLOR\n",
+        "    bg_image = np.zeros(image_data.shape, dtype=np.uint8)\n",
+        "    bg_image[:] = BG_COLOR\n",
+        "\n",
+        "    condition = np.stack((category_masks[0].numpy_view(),) * 3, axis=-1) > 0.2\n",
+        "    output_image = np.where(condition, fg_image, bg_image)\n",
+        "\n",
+        "    print(f'Segmentation mask of {name}:')\n",
+        "    resize_and_show(output_image)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "Now that you know how to separate the foreground and background of an image, let's take it a step further and blur the background for an effect similar to what is provided by Google Hangouts."
+      ],
+      "metadata": {
+        "id": "-RxsX2CjJbfo"
+      }
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ciCGwCQ3gNDc"
+      },
+      "outputs": [],
+      "source": [
+        "# Blur the image background based on the segmentation mask.\n",
+        "\n",
+        "# Create the segmenter\n",
+        "with python.vision.ImageSegmenter.create_from_options(options) as segmenter:\n",
+        "\n",
+        "  # Loop through available image(s)\n",
+        "  for image_file_name in IMAGE_FILENAMES:\n",
+        "\n",
+        "    # Create the MediaPipe Image\n",
+        "    image = mp.Image.create_from_file(image_file_name)\n",
+        "\n",
+        "    # Retrieve the category masks for the image\n",
+        "    category_masks = segmenter.segment(image)\n",
+        "\n",
+        "    # Convert the BGR image to RGB\n",
+        "    image_data = cv2.cvtColor(image.numpy_view(), cv2.COLOR_BGR2RGB)\n",
+        "\n",
+        "    # Apply effects\n",
+        "    blurred_image = cv2.GaussianBlur(image_data, (55,55), 0)\n",
+        "    condition = np.stack((category_masks[0].numpy_view(),) * 3, axis=-1) > 0.1\n",
+        "    output_image = np.where(condition, image_data, blurred_image)\n",
+        "\n",
+        "    print(f'Blurred background of {image_file_name}:')\n",
+        "    resize_and_show(output_image)"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "provenance": []
+    },
+    "kernelspec": {
+      "display_name": "Python 3 (ipykernel)",
+      "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.7.13"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}