notebooks/how-to-generate-segmentation-mask-with-roboflow.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "#@title String fields\n",
        "\n",
        "path_to_image = './PATH_TO_IMAGE' #@param {type:\"string\"}\n",
        "api_key = 'YOUR_API_KEY' #@param {type:\"string\"}\n",
        "model = 'YOUR_MODEL_NAME' #@param {type:\"string\"}\n",
        "version = 'YOUR_VERSION_NUMBER' #@param {type:\"string\"}"
      ],
      "metadata": {
        "id": "22fP_uGRgzww"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!pip3 install requests_toolbelt"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "yDq9PNhHXlUJ",
        "outputId": "1864d8d5-70b7-4c2f-c5d8-591a13c3985f"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting requests_toolbelt\n",
            "  Downloading requests_toolbelt-0.10.1-py2.py3-none-any.whl (54 kB)\n",
            "\u001b[K     |████████████████████████████████| 54 kB 296 kB/s \n",
            "\u001b[?25hRequirement already satisfied: requests<3.0.0,>=2.0.1 in /usr/local/lib/python3.7/dist-packages (from requests_toolbelt) (2.23.0)\n",
            "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from requests<3.0.0,>=2.0.1->requests_toolbelt) (1.24.3)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests<3.0.0,>=2.0.1->requests_toolbelt) (2022.9.24)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests<3.0.0,>=2.0.1->requests_toolbelt) (2.10)\n",
            "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests<3.0.0,>=2.0.1->requests_toolbelt) (3.0.4)\n",
            "Installing collected packages: requests-toolbelt\n",
            "Successfully installed requests-toolbelt-0.10.1\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy, io\n",
        "from PIL import Image, ImageDraw\n",
        "import cv2\n",
        "import requests\n",
        "from requests_toolbelt.multipart.encoder import MultipartEncoder"
      ],
      "metadata": {
        "id": "PCCrKGR0BpEh"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Helper Functions"
      ],
      "metadata": {
        "id": "jZhL3Ni0hI0-"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def infer_with_segmentation(image_path, api_key, model, version):\n",
        "    # Load Image with PIL\n",
        "    open_img = cv2.imread(image_path)\n",
        "    color_img = cv2.cvtColor(open_img, cv2.COLOR_BGR2RGB)\n",
        "    pil_image = Image.fromarray(color_img)\n",
        "\n",
        "    # Convert to JPEG Buffer\n",
        "    buffered = io.BytesIO()\n",
        "    pil_image.save(buffered, quality=100, format=\"JPEG\")\n",
        "\n",
        "    # Build multipart form and post request\n",
        "    m = MultipartEncoder(fields={'file': (\"imageToUpload\", buffered.getvalue(), \"image/jpeg\")})\n",
        "\n",
        "    response = requests.post(f\"https://outline.roboflow.com/{model}/{version}?api_key={api_key}\", data=m, headers={'Content-Type': m.content_type})\n",
        "\n",
        "    # formatted = json.dumps(response.json(), indent=2)\n",
        "    # print(formatted)\n",
        "\n",
        "    return response.json()"
      ],
      "metadata": {
        "id": "VkWRAi7Xgw5J"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Main Code"
      ],
      "metadata": {
        "id": "4pqPGPN2hMcs"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# perform inference on the image\n",
        "results = infer_with_segmentation(path_to_image, api_key, model, version)\n",
        "\n",
        "# convert image to numpy (for convenience)\n",
        "imArray = numpy.asarray(Image.open(path_to_image).convert(\"RGBA\"))\n",
        "\n",
        "# process each resulting prediction \n",
        "for index, prediction in enumerate(results[\"predictions\"]):\n",
        "\n",
        "  # convert polygon prediction points to workable tuple format\n",
        "  polygon = [(point_obj[\"x\"], point_obj[\"y\"]) for point_obj in prediction[\"points\"]]\n",
        "\n",
        "  # create a mask image\n",
        "  maskIm = Image.new('L', (imArray.shape[1], imArray.shape[0]), 0)\n",
        "  ImageDraw.Draw(maskIm).polygon(polygon, outline=1, fill=1)\n",
        "\n",
        "  # convert to numpy array\n",
        "  maskArray = numpy.array(maskIm)\n",
        "\n",
        "  # assemble new image (uint8: 0-255)\n",
        "  newImArray = numpy.empty(imArray.shape,dtype='uint8')\n",
        "\n",
        "  # colors (three first columns, RGB)\n",
        "  newImArray[:,:,:3] = imArray[:,:,:3]\n",
        "\n",
        "  # transparency (4th column)\n",
        "  newImArray[:,:,3] = maskArray*255\n",
        "\n",
        "  # convert back to Image from numpy and save\n",
        "  newIm = Image.fromarray(newImArray, \"RGBA\")\n",
        "  newIm.save(f\"mask_{index}.png\")"
      ],
      "metadata": {
        "id": "jIilHjUsBkVF"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}