MODNet_Image_Matting_Demo.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/fixed0301/OpenEYE-my/blob/master/MODNet_Image_Matting_Demo.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wlARSYHGHgev"
      },
      "source": [
        "# MODNet - Portrait Image Matting Demo\n",
        "\n",
        "[![arXiv](https://img.shields.io/badge/arXiv-Paper-<COLOR>.svg)](https://arxiv.org/pdf/2011.11961.pdf)          [![GitHub stars](https://img.shields.io/github/stars/ZHKKKe/MODNet?style=social)](https://github.com/ZHKKKe/MODNet)\n",
        "\n",
        "<p align=\"justify\">This is a <b>Portrait Image Matting Demo</b> of our paper ''<a href=\"https://arxiv.org/pdf/2011.11961.pdf\">Is a Green Screen Really Necessary for Real-Time Portrait Matting?</a>''. We propose a trimap-free MODNet for portrait matting in real time (on a single GPU). If you want to try portrait video matting with MODNet, please refer to our <b>Real-Time Portrait Video Matting Demo</b> based on WebCam [<a href=\"https://colab.research.google.com/drive/1Pt3KDSc2q7WxFvekCnCLD8P0gBEbxm6J?usp=sharing\">online demo</a>] [<a href=\"https://github.com/ZHKKKe/MODNet/tree/master/demo/video_matting\">offline demo</a>].</p>\n",
        "\n",
        "<img src=\"https://raw.githubusercontent.com/ZHKKKe/MODNet/develop/doc/gif/image_matting_demo.gif\" width=\"40%\">\n",
        "\n",
        "<p align=\"justify\">In this demo, we provide a pre-trained MODNet for photographic portrait images (<i>e.g.</i>, the image from <a href=\"https://www.flickr.com/\">Flickr</a>). Of course, this model can also handle the portrait images from other data domains (<i>e.g.</i>, the image captured by smartphone). Note that <b><font color='#FF000'>our model may still fail in some cases (<i>e.g.</i>, complex backgrounds and challenging poses/costumes) due to our relatively small labeled training set (only 3k labeled foregrounds)</font></b>.</p>\n",
        "\n",
        "### **Let's start!**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F3t-jsyrtl7w"
      },
      "source": [
        "## 1. Preparation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jkk60QsJw8D0"
      },
      "source": [
        "<p align=\"justify\">In the top menu of this session, select <b>Runtime -> Change runtime type</b>, and set <b>Hardware Accelerator</b> to <b>GPU</b>.</p>\n",
        "\n",
        "<p align=\"justify\">Clone the repository, and download the pre-trained model:</p>\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gEwiokUgkSDO"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "# clone the repository\n",
        "%cd /content\n",
        "if not os.path.exists('MODNet'):\n",
        "  !git clone https://github.com/ZHKKKe/MODNet\n",
        "%cd MODNet/\n",
        "\n",
        "# dowload the pre-trained ckpt for image matting\n",
        "pretrained_ckpt = 'pretrained/modnet_photographic_portrait_matting.ckpt'\n",
        "if not os.path.exists(pretrained_ckpt):\n",
        "  !gdown --id 1mcr7ALciuAsHCpLnrtG_eop5-EYhbCmz \\\n",
        "          -O pretrained/modnet_photographic_portrait_matting.ckpt\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yaXFIdEpulEi"
      },
      "source": [
        "## 2. Upload Images"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Nog7bf3nfG6Y"
      },
      "source": [
        "<p align=\"justify\">Upload portrait images to be processed (only PNG and JPG format are supported):</p>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c_3jhVwQu3-P"
      },
      "outputs": [],
      "source": [
        "import shutil\n",
        "from google.colab import files\n",
        "\n",
        "# clean and rebuild the image folders\n",
        "input_folder = 'demo/image_matting/colab/input'\n",
        "if os.path.exists(input_folder):\n",
        "  shutil.rmtree(input_folder)\n",
        "os.makedirs(input_folder)\n",
        "\n",
        "output_folder = 'demo/image_matting/colab/output'\n",
        "if os.path.exists(output_folder):\n",
        "  shutil.rmtree(output_folder)\n",
        "os.makedirs(output_folder)\n",
        "\n",
        "# upload images (PNG or JPG)\n",
        "image_names = list(files.upload().keys())\n",
        "for image_name in image_names:\n",
        "  shutil.move(image_name, os.path.join(input_folder, image_name))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wwMfEIiwu5VB"
      },
      "source": [
        "## 3. Inference"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jJaRhVYdfxNt"
      },
      "source": [
        "<p align=\"justify\">Run the following command for alpha matte prediction:</p>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fi1jtFug-Nvh"
      },
      "outputs": [],
      "source": [
        "!python -m demo.image_matting.colab.inference \\\n",
        "        --input-path demo/image_matting/colab/input \\\n",
        "        --output-path demo/image_matting/colab/output \\\n",
        "        --ckpt-path ./pretrained/modnet_photographic_portrait_matting.ckpt\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CwGwwojKu_en"
      },
      "source": [
        "## 4. Visualization"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u6FP75C8gHbz"
      },
      "source": [
        "<p align=\"justify\">Display the results (from left to right: image, foreground, and alpha matte):</p>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JOmYOHKfgQ5Y"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "from PIL import Image\n",
        "\n",
        "def combined_display(image, matte):\n",
        "  # calculate display resolution\n",
        "  w, h = image.width, image.height\n",
        "  rw, rh = 800, int(h * 800 / (3 * w))\n",
        "\n",
        "  # obtain predicted foreground\n",
        "  image = np.asarray(image)\n",
        "  if len(image.shape) == 2:\n",
        "    image = image[:, :, None]\n",
        "  if image.shape[2] == 1:\n",
        "    image = np.repeat(image, 3, axis=2)\n",
        "  elif image.shape[2] == 4:\n",
        "    image = image[:, :, 0:3]\n",
        "  matte = np.repeat(np.asarray(matte)[:, :, None], 3, axis=2) / 255\n",
        "  foreground = image * matte + np.full(image.shape, 255) * (1 - matte)\n",
        "\n",
        "  # combine image, foreground, and alpha into one line\n",
        "  combined = np.concatenate((image, foreground, matte * 255), axis=1)\n",
        "  combined = Image.fromarray(np.uint8(combined)).resize((rw, rh))\n",
        "  return combined\n",
        "\n",
        "# visualize all images\n",
        "image_names = os.listdir(input_folder)\n",
        "for image_name in image_names:\n",
        "  matte_name = image_name.split('.')[0] + '.png'\n",
        "  image = Image.open(os.path.join(input_folder, image_name))\n",
        "  matte = Image.open(os.path.join(output_folder, matte_name))\n",
        "  display(combined_display(image, matte))\n",
        "  print(image_name, '\\n')\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yfDKNQNeeLhx"
      },
      "source": [
        "## Download Results"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TqNdGeIf1cc8"
      },
      "source": [
        "<p align=\"justify\">Download the Zip package of predicted alpha mattes:</p>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e6aqVu631k01"
      },
      "outputs": [],
      "source": [
        "zip_filename = 'matte.zip'\n",
        "if os.path.exists(zip_filename):\n",
        "  os.remove(zip_filename)\n",
        "\n",
        "os.system(f\"zip -r -j {zip_filename} {output_folder}/*\")\n",
        "files.download(zip_filename)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}