rembg.ipynb

rembg.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "rembg.ipynb",
      "provenance": [],
      "authorship_tag": "ABX9TyPC2QJ1POv+P/QVnV4zIdFj",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/gist/yui0/16f1c560fa2763d66d1103d9191ee5c4/rembg.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2Dyscyd4TbCT"
      },
      "source": [
        "#!pip install torch==1.7.1+cpu torchvision==0.8.2+cpu -f https://download.pytorch.org/whl/torch_stable.html"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sSwv98gJaEEH"
      },
      "source": [
        "!pip install rembg"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3CpzDHb7bxXY"
      },
      "source": [
        "!cat /usr/local/lib/python3.7/dist-packages/rembg/bg.py"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QaULvZtBYJWz"
      },
      "source": [
        "!sed -e 's/@functools.lru_cache/@functools.lru_cache(maxsize=None)/' -i.bak /usr/local/lib/python3.7/dist-packages/rembg/bg.py\n",
        "!sed -e 's:mask = detect.predict(model, np.array(img)).convert(\"L\"):mask = detect.predict(model, np.array(img)).convert(\"L\")\\n    output_path = \"/content/trimap.png\"\\n    mask = mask.resize(img.size)\\n    mask.save(output_path):' -i.bak /usr/local/lib/python3.7/dist-packages/rembg/bg.py"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "a2B4Oucuae0p"
      },
      "source": [
        "!cat /usr/local/lib/python3.7/dist-packages/rembg/bg.py"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Nw-2VydCUBMV"
      },
      "source": [
        "from google.colab import files\n",
        "uploaded = files.upload()\n",
        "name = list(uploaded.keys())[0]\n",
        "\n",
        "name"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-NZcYIFUWe9Z"
      },
      "source": [
        "from rembg.bg import remove\n",
        "import numpy as np\n",
        "import io\n",
        "from PIL import Image\n",
        "\n",
        "from PIL import ImageFile\n",
        "ImageFile.LOAD_TRUNCATED_IMAGES = True\n",
        "\n",
        "input_path = name\n",
        "output_path = 'rembg.png'\n",
        "\n",
        "f = np.fromfile(input_path)\n",
        "result = remove(f)\n",
        "img = Image.open(io.BytesIO(result)).convert(\"RGBA\")\n",
        "img.save(output_path)"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}