mediapipe_slimnet/desktop/mediapipe_custom_calculator.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "mediapipe_custom_calculator.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TOhzfnL3Dwp5",
        "colab_type": "text"
      },
      "source": [
        "**Portrait Video Segmentation Using Custom Medaipipe Calculator**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4e6dzoC5CFj9",
        "colab_type": "text"
      },
      "source": [
        "In this demo, we will build a portrait segmentaion aplication using **custom calculators** on **desktop**, using mediapipe. There will be two video file inputs and one video fiile output. Our aim is to **blend** the portrait foreground region into the background video, with the help of segmentaion **mask**. As in the case of android, we will follow the basic segmentation pipeline from **hair segmentaion** example. Since the application uses **gpu** operations, choose a GPU runtime for development and deployment."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "m6XRrTgPgCgV",
        "colab_type": "text"
      },
      "source": [
        "**1. Checkout MediaPipe Github Repository**\n",
        "\n",
        "The mediapipe **repository** contains many demo applications for android. We will modify the hair_segmentaion application, which contains the basic pipeline for **video segmentaion**."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0PQ0WAmffiI_",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!git clone https://github.com/google/mediapipe.git\n",
        "%cd mediapipe\n",
        "!sudo apt install curl gnupg\n",
        "!curl -fsSL https://bazel.build/bazel-release.pub.gpg | gpg --dearmor > bazel.gpg\n",
        "!sudo mv bazel.gpg /etc/apt/trusted.gpg.d/\n",
        "!echo \"deb [arch=amd64] https://storage.googleapis.com/bazel-apt stable jdk1.8\" | sudo tee /etc/apt/sources.list.d/bazel.list\n",
        "!sudo apt update && sudo apt install bazel"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hou9qtIzhDPG",
        "colab_type": "text"
      },
      "source": [
        "**3. Install a JDK (optional)**\n",
        "\n",
        "Sometimes the **default jdk** version may cause an **error** during android sdk installation in ubuntu. So, install an older version of **openjdk-8** and configure the same as the default version of the system."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2AAP-r1ogbSp",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!sudo apt install openjdk-8-jdk\n",
        "!sudo update-alternatives --config java  # Choose OpenJDK 8\n",
        "!java -version"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "duUg0dBmVl5Q",
        "colab_type": "text"
      },
      "source": [
        "**3. Install OpenCV (optional)**\n",
        "\n",
        "If opencv is not already installed, run **setup_opencv.sh** to automatically build OpenCV from source and modify MediaPipe’s OpenCV config."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OOUYEIwDWIaB",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!bash setup_opencv.sh"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qnmfyimNXQlm",
        "colab_type": "text"
      },
      "source": [
        "**4. Install MediaPipe Without Android Studio (SDK & NDK)**\n",
        "\n",
        "If the android studio is not installed in your system, you can configute mediapipe with **sdk and ndk** by running this script."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DK1_chxlXbhQ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!bash setup_android_sdk_and_ndk.sh"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "z59JKpw0XjgY",
        "colab_type": "text"
      },
      "source": [
        "**Note:** If Android SDK and NDK are already installed, set **ANDROID_HOME** and **ANDROID_NDK_HOME** paths accordingly.\n",
        "\n",
        "```\n",
        "export ANDROID_HOME=<path to the Android SDK>\n",
        "export ANDROID_NDK_HOME=<path to the Android NDK>\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "maPLx3SbZtqn",
        "colab_type": "text"
      },
      "source": [
        "**6. Modify The Hair Segmentaion Mediapipe Application**\n",
        "\n",
        "Initially put the **[portrait_segmentation](https://github.com/anilsathyan7/Portrait-Segmentation/blob/master/mediapipe_slimnet/portrait_segmentation.tflite)** tflite file into **models** directory, inside mediapipe folder."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4CMvQZZbaVpi",
        "colab_type": "text"
      },
      "source": [
        "**A.** Create a new directory called **portrait_segmentation** under **graphs** subdirectory and copy all the files from **hair_segmentation**."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Q9VOETHjaU4E",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!mkdir mediapipe/graphs/portrait_segmentation\n",
        "!cp -r mediapipe/graphs/hair_segmentation/* mediapipe/graphs/portrait_segmentation\n",
        "!mv mediapipe/graphs/portrait_segmentation/hair_segmentation_mobile_gpu.pbtxt mediapipe/graphs/portrait_segmentation/portrait_segmentation_mobile_gpu.pbtxt"
      ],
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c1g8bpWRakmg",
        "colab_type": "text"
      },
      "source": [
        "Rename the pbtxt file(mv command) as **portrait_segmentation_mobile_gpu.pbtxt** and modify the following lines :-\n",
        "\n",
        "1. **Number of channels**: 'max_num_channels: 4' in **TfLiteConverterCalculator** to max_num_channels: 3\n",
        "2. **Model name**: 'hair_segmentaion.tflite' in **TfLiteInferenceCalculator** to portrait_segmentation.tflite\n",
        "\n",
        "\n",
        "**B**. Add two new nodes for video file input i.e **OpenCvVideoDecoderCalculator** and one node for video file output i.e **OpenCvVideoEncoderCalculator** in the new pipeline.\n",
        "\n",
        "```\n",
        "# Decodes an input video file into images and a video header.\n",
        "node {\n",
        "  calculator: \"OpenCvVideoDecoderCalculator\"\n",
        "  input_side_packet: \"INPUT_FILE_PATH:input_video_path\"\n",
        "  output_stream: \"VIDEO:input_video\"\n",
        "  output_stream: \"VIDEO_PRESTREAM:input_video_header\"\n",
        "}\n",
        "\n",
        "# Decodes an input video file into images and a video header.\n",
        "node {\n",
        "  calculator: \"OpenCvVideoDecoderCalculator\"\n",
        "  input_side_packet: \"INPUT_FILE_PATH:side_video_path\"\n",
        "  output_stream: \"VIDEO:side_video\"\n",
        "  output_stream: \"VIDEO_PRESTREAM:side_video_header\"\n",
        "}\n",
        "\n",
        "# Encodes the annotated images into a video file, adopting properties specified\n",
        "# in the input video header, e.g., video framerate.\n",
        "node {\n",
        "  calculator: \"OpenCvVideoEncoderCalculator\"\n",
        "  input_stream: \"VIDEO:output_video\"\n",
        "  input_stream: \"VIDEO_PRESTREAM:input_video_header\"\n",
        "  input_side_packet: \"OUTPUT_FILE_PATH:output_video_path\"\n",
        "  node_options: {\n",
        "    [type.googleapis.com/mediapipe.OpenCvVideoEncoderCalculatorOptions]: {\n",
        "      codec: \"avc1\"\n",
        "      video_format: \"mp4\"\n",
        "    }\n",
        "  }\n",
        "}\n",
        "```\n",
        "\n",
        "**Note:** The main input should be a video file contaning portrait images and the other one should be a background video.\n",
        "\n",
        "**C.** Now remove the 'RecolorCalculator' node, and instead add the custom **SeamlessCloningCalculator**  into the pipeline.\n",
        "\n",
        "```\n",
        "# Takes Image, Mask and Background as input and performs \n",
        "# poisson blending using opencv library\n",
        "node {\n",
        "  calculator: \"SeamlessCloningCalculator\"\n",
        "  input_stream: \"IMAGE_CPU:input_video\"\n",
        "  input_stream: \"BACKGROUND_CPU:sync_side_video\"\n",
        "  input_stream: \"MASK_CPU:portrait_mask_cpu\"\n",
        "  output_stream: \"OUTPUT_VIDEO:output_video\"\n",
        "}\n",
        "```\n",
        "**Note:** The idea is to combine the foreground in the image with background using mask, such that portrait foreground blends into the background image seamlessly.\n",
        "\n",
        "**D**. Use **PacketClonerCalculator** to clone the background video frames when all frames are used up. Also use **ImageFrameToGpuBufferCalculator** and **GpuBufferToImageFrameCalculator** for copying data between CPU and GPU, whenever necessary.\n",
        "\n",
        "**E**. Now, inside the **BUILD** file in this directory(portrait_segmentation), change the graph name to \"**portrait_segmentation_mobile_gpu.pbtxt**\".\n",
        "\n",
        "Also add the  **calculator files** inside the **cc_library** section for mobile_calcualators as follows:-\n",
        "\n",
        "```\n",
        "\"//mediapipe/calculators/video:opencv_video_decoder_calculator\",\n",
        "\"//mediapipe/calculators/video:opencv_video_encoder_calculator\",\n",
        "\"//mediapipe/calculators/image:poisson_blending_calculator\",\n",
        "\"//mediapipe/calculators/core:packet_cloner_calculator\",\n",
        "```\n",
        "\n",
        "See the final pbtxt file: [portrait_segmentation_mobile_gpu.pbtxt](https://github.com/anilsathyan7/Portrait-Segmentation/blob/master/mediapipe_slimnet/desktop/portrait_segmentation_mobile_gpu.pbtxt)\n",
        "\n",
        "\n",
        "**F**. Similary, create a new folder called '**portrait_segmentation**' inside example directory at location: '**/mediapipe/examples/desktop/**'.Add the BUILD file  inside this directory as shown below.\n",
        "\n",
        "```\n",
        "licenses([\"notice\"])\n",
        "package(default_visibility = [\"//mediapipe/examples:__subpackages__\"])\n",
        "\n",
        "# Linux only\n",
        "cc_binary(\n",
        "    name = \"portrait_segmentation_gpu\",\n",
        "    deps = [\n",
        "        \"//mediapipe/examples/desktop:simple_run_graph_main\",\n",
        "        \"//mediapipe/examples/desktop:demo_run_graph_main_gpu\",\n",
        "        \"//mediapipe/graphs/portrait_segmentation:mobile_calculators\",\n",
        "    ],\n",
        ")\n",
        "```\n",
        "**Note:** We added a dependency '**simple_run_graph_main**' for executing graph using side packets and video file inputs.\n",
        "\n",
        "**G.** Add the following dependencies into the **BUILD** file inside **calculators/image** directory.\n",
        "\n",
        "```\n",
        "cc_library(\n",
        "    name = \"poisson_blending_calculator\",\n",
        "    srcs = [\"poisson_blending_calculator.cc\"],\n",
        "    visibility = [\"//visibility:public\"],    \n",
        "    deps = [\n",
        "        \"//mediapipe/gpu:gl_calculator_helper\",\n",
        "        \"//mediapipe/framework:calculator_framework\",\n",
        "        \"//mediapipe/framework:calculator_options_cc_proto\",\n",
        "        \"//mediapipe/framework:timestamp\",\n",
        "        \"//mediapipe/framework/port:status\",\n",
        "        \"//mediapipe/framework/deps:file_path\",\n",
        "        \"@com_google_absl//absl/time\",\n",
        "        \"@com_google_absl//absl/strings\",\n",
        "        \"//mediapipe/framework/formats:rect_cc_proto\",\n",
        "        \"//mediapipe/framework/port:ret_check\",\n",
        "        \"//mediapipe/framework/formats:image_frame\",\n",
        "        \"//mediapipe/framework/formats:matrix\",\n",
        "        \"//mediapipe/framework/formats:image_frame_opencv\",\n",
        "        \"//mediapipe/framework/port:opencv_core\",\n",
        "        \"//mediapipe/framework/port:opencv_imgproc\",\n",
        "        \"//mediapipe/framework/port:opencv_imgcodecs\", \n",
        "        \"//mediapipe/util:resource_util\",\n",
        "        ],\n",
        "    alwayslink = 1,\n",
        ")\n",
        "```\n",
        "**Note:** The file **poisson_blending_calculator** refers to our custom seamlesscloning calculator C++ file.\n",
        "\n",
        "**H.** Add the following lines under cc_library section in the **opencv_linux.BUILD** file, inside the third_party directory.\n",
        "\n",
        "`\"lib/libopencv_photo.so\",`\n",
        "or\n",
        "`\"lib/x86_64-linux-gnu/libopencv_photo.so\",`\n",
        "\n",
        "**Note:** Make sure the file exists at the specified path in the system. This ensures that opencv can link the files for seamlesscloning from the module **photo** during build."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6jX_IrmM19_Q",
        "colab_type": "text"
      },
      "source": [
        "**Seamless Clone Custom Calculator**\n",
        "\n",
        "A custom calculator can created by defining a new sub-class of the CalculatorBase class, implementing a number of methods, and registering the new sub-class with Mediapipe. At a minimum, a new calculator must implement the following methods - **GetContract, Open, Process and Close.**\n",
        "\n",
        "The **Process** method continously takes inputs, processes them and produce outputs. We will write most of our code for seamless cloning within this method; whereas in **GetContract** we just specify the expected types of inputs and outputs of the calculator."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HKoKbmQV2Gr_",
        "colab_type": "text"
      },
      "source": [
        "*Steps for seamless cloning:-*\n",
        "\n",
        "\n",
        "1.  Convert the mask to binary Mat format with 0's representing background and 1 foreground.\n",
        "\n",
        "2.  Resize the mask and background image to the size of the input image.\n",
        "\n",
        "3.  Dilate the mask to include neighbouring background regions around borders.\n",
        "\n",
        "4.  Find the largest contour and corresponding bounding rectangle from mask.\n",
        "\n",
        "5.  Crop out the ROI from input and mask image using the bounding rectangle.\n",
        "\n",
        "6.  Set the foreground pixels values of the mask to 255.\n",
        "\n",
        "7.  Calculate the location of the center of the input roi image in the background.\n",
        "\n",
        "8.  Perform seamless cloning of input on background, using mask and return the result.\n",
        "\n",
        "Thus, the seamless clone calculator takes **three inputs** as CPU ImageFrame's i.e input image, background image and segmentaion mask. It produces a single ouput image frame,  representing the **blended image** in CPU. Finally, we save the results into a video file using **OpenCvVideoEncoderCalculator**.\n",
        "\n",
        "Copy the calculator file '**[poisson_blending_calculator.cc](https://github.com/anilsathyan7/Portrait-Segmentation/blob/master/mediapipe_slimnet/desktop/poisson_blending_calculator.cc)**' into the directory - **mediapipe/calculators/image**.\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hOkkPlkHW5Rn",
        "colab_type": "text"
      },
      "source": [
        "**Build and Run**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cLeeSDzdK1t5",
        "colab_type": "text"
      },
      "source": [
        "To build the portrait_segmentation_gpu app on **desktop** using bazel, run:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kUusPEByW896",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!bazel build -c opt --copt -DMESA_EGL_NO_X11_HEADERS --copt -DEGL_NO_X11 mediapipe/examples/desktop/portrait_segmentation:portrait_segmentation_gpu"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UYkvtIoILGnA",
        "colab_type": "text"
      },
      "source": [
        "Now, load two video files (i.e **portrait** and **background**) as inputs, run the application and save the output video."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ba2BeAeqkkqi",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!GLOG_logtostderr=1 bazel-bin/mediapipe/examples/desktop/portrait_segmentation/portrait_segmentation_gpu --calculator_graph_config_file=/content/mediapipe/mediapipe/graphs/portrait_segmentation/portrait_segmentation_mobile_gpu.pbtxt --input_side_packets=side_video_path=/content/fire_vid.mp4,input_video_path=/content/grandma_vid.mp4,output_video_path=/content/output15.mp4"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7-nWK9Sh1gdN",
        "colab_type": "text"
      },
      "source": [
        "**Note:** If the run failed in a headless set-up(eg. google colab), then modify the file mediapipe/gpu/**gl_context_egl.cc** by removing the option '**EGL_WINDOW_BIT** and rebuild the application."
      ]
    }
  ]
}