People_segmentation.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "“People segmentation1”的副本",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/lkp520/peo_seg/blob/main/%E2%80%9CPeople_segmentation1%E2%80%9D%E7%9A%84%E5%89%AF%E6%9C%AC.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5fVuyHf6zstb"
      },
      "source": [
        "%matplotlib inline"
      ],
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qPfDG7jKz1FF"
      },
      "source": [
        "from pylab import imshow"
      ],
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "a_goHaGlBjZf"
      },
      "source": [
        "import numpy as np"
      ],
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yx4oLfemz4ys"
      },
      "source": [
        "import os\n",
        "import cv2"
      ],
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Mp2G_NgPBByg"
      },
      "source": [
        "import torch"
      ],
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tC_aZXqr_4np"
      },
      "source": [
        "import albumentations as albu"
      ],
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "K6y5PqbOprr6",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "11b2ef32-eab8-475d-c582-6ed2951875a0"
      },
      "source": [
        "!pip install iglovikov_helper_functions"
      ],
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Requirement already satisfied: iglovikov_helper_functions in /usr/local/lib/python3.7/dist-packages (0.0.53)\n",
            "Requirement already satisfied: joblib in /usr/local/lib/python3.7/dist-packages (from iglovikov_helper_functions) (1.0.1)\n",
            "Requirement already satisfied: scikit-learn in /usr/local/lib/python3.7/dist-packages (from iglovikov_helper_functions) (0.22.2.post1)\n",
            "Requirement already satisfied: scikit-image in /usr/local/lib/python3.7/dist-packages (from iglovikov_helper_functions) (0.16.2)\n",
            "Requirement already satisfied: pandas in /usr/local/lib/python3.7/dist-packages (from iglovikov_helper_functions) (1.1.5)\n",
            "Requirement already satisfied: jpeg4py in /usr/local/lib/python3.7/dist-packages (from iglovikov_helper_functions) (0.1.4)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from iglovikov_helper_functions) (1.19.5)\n",
            "Requirement already satisfied: scipy in /usr/local/lib/python3.7/dist-packages (from iglovikov_helper_functions) (1.4.1)\n",
            "Requirement already satisfied: Pillow in /usr/local/lib/python3.7/dist-packages (from iglovikov_helper_functions) (7.1.2)\n",
            "Requirement already satisfied: imagecorruptions in /usr/local/lib/python3.7/dist-packages (from iglovikov_helper_functions) (1.1.2)\n",
            "Requirement already satisfied: torch in /usr/local/lib/python3.7/dist-packages (from iglovikov_helper_functions) (1.9.0+cu102)\n",
            "Requirement already satisfied: addict in /usr/local/lib/python3.7/dist-packages (from iglovikov_helper_functions) (2.4.0)\n",
            "Requirement already satisfied: opencv-python in /usr/local/lib/python3.7/dist-packages (from iglovikov_helper_functions) (4.1.2.30)\n",
            "Requirement already satisfied: matplotlib!=3.0.0,>=2.0.0 in /usr/local/lib/python3.7/dist-packages (from scikit-image->iglovikov_helper_functions) (3.2.2)\n",
            "Requirement already satisfied: networkx>=2.0 in /usr/local/lib/python3.7/dist-packages (from scikit-image->iglovikov_helper_functions) (2.5.1)\n",
            "Requirement already satisfied: imageio>=2.3.0 in /usr/local/lib/python3.7/dist-packages (from scikit-image->iglovikov_helper_functions) (2.4.1)\n",
            "Requirement already satisfied: PyWavelets>=0.4.0 in /usr/local/lib/python3.7/dist-packages (from scikit-image->iglovikov_helper_functions) (1.1.1)\n",
            "Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.7/dist-packages (from matplotlib!=3.0.0,>=2.0.0->scikit-image->iglovikov_helper_functions) (0.10.0)\n",
            "Requirement already satisfied: python-dateutil>=2.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib!=3.0.0,>=2.0.0->scikit-image->iglovikov_helper_functions) (2.8.1)\n",
            "Requirement already satisfied: pyparsing!=2.0.4,!=2.1.2,!=2.1.6,>=2.0.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib!=3.0.0,>=2.0.0->scikit-image->iglovikov_helper_functions) (2.4.7)\n",
            "Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib!=3.0.0,>=2.0.0->scikit-image->iglovikov_helper_functions) (1.3.1)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.7/dist-packages (from cycler>=0.10->matplotlib!=3.0.0,>=2.0.0->scikit-image->iglovikov_helper_functions) (1.15.0)\n",
            "Requirement already satisfied: decorator<5,>=4.3 in /usr/local/lib/python3.7/dist-packages (from networkx>=2.0->scikit-image->iglovikov_helper_functions) (4.4.2)\n",
            "Requirement already satisfied: cffi in /usr/local/lib/python3.7/dist-packages (from jpeg4py->iglovikov_helper_functions) (1.14.6)\n",
            "Requirement already satisfied: pycparser in /usr/local/lib/python3.7/dist-packages (from cffi->jpeg4py->iglovikov_helper_functions) (2.20)\n",
            "Requirement already satisfied: pytz>=2017.2 in /usr/local/lib/python3.7/dist-packages (from pandas->iglovikov_helper_functions) (2018.9)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from torch->iglovikov_helper_functions) (3.7.4.3)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FgZbmE9f_Vf-"
      },
      "source": [
        "from iglovikov_helper_functions.utils.image_utils import load_rgb, pad, unpad\n",
        "from iglovikov_helper_functions.dl.pytorch.utils import tensor_from_rgb_image"
      ],
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NluTQDch1Z0X"
      },
      "source": [
        "# !wget https://habrastorage.org/webt/em/l7/cr/eml7crxnxftrimsmolwjegqcrp4.jpeg > /dev/null"
      ],
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6-Y7GQabBj3C",
        "outputId": "469e08bc-9ce0-4974-de95-6b65617c01f0"
      },
      "source": [
        "!pip install tensorboard==2.5"
      ],
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Collecting tensorboard==2.5\n",
            "  Downloading tensorboard-2.5.0-py3-none-any.whl (6.0 MB)\n",
            "\u001b[K     |████████████████████████████████| 6.0 MB 5.6 MB/s \n",
            "\u001b[?25hRequirement already satisfied: markdown>=2.6.8 in /usr/local/lib/python3.7/dist-packages (from tensorboard==2.5) (3.3.4)\n",
            "Requirement already satisfied: google-auth<2,>=1.6.3 in /usr/local/lib/python3.7/dist-packages (from tensorboard==2.5) (1.32.1)\n",
            "Requirement already satisfied: tensorboard-data-server<0.7.0,>=0.6.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard==2.5) (0.6.1)\n",
            "Requirement already satisfied: setuptools>=41.0.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard==2.5) (57.2.0)\n",
            "Requirement already satisfied: tensorboard-plugin-wit>=1.6.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard==2.5) (1.8.0)\n",
            "Requirement already satisfied: werkzeug>=0.11.15 in /usr/local/lib/python3.7/dist-packages (from tensorboard==2.5) (1.0.1)\n",
            "Requirement already satisfied: google-auth-oauthlib<0.5,>=0.4.1 in /usr/local/lib/python3.7/dist-packages (from tensorboard==2.5) (0.4.4)\n",
            "Requirement already satisfied: requests<3,>=2.21.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard==2.5) (2.23.0)\n",
            "Requirement already satisfied: grpcio>=1.24.3 in /usr/local/lib/python3.7/dist-packages (from tensorboard==2.5) (1.34.1)\n",
            "Requirement already satisfied: wheel>=0.26 in /usr/local/lib/python3.7/dist-packages (from tensorboard==2.5) (0.36.2)\n",
            "Requirement already satisfied: protobuf>=3.6.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard==2.5) (3.17.3)\n",
            "Requirement already satisfied: absl-py>=0.4 in /usr/local/lib/python3.7/dist-packages (from tensorboard==2.5) (0.12.0)\n",
            "Requirement already satisfied: numpy>=1.12.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard==2.5) (1.19.5)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.7/dist-packages (from absl-py>=0.4->tensorboard==2.5) (1.15.0)\n",
            "Requirement already satisfied: cachetools<5.0,>=2.0.0 in /usr/local/lib/python3.7/dist-packages (from google-auth<2,>=1.6.3->tensorboard==2.5) (4.2.2)\n",
            "Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/local/lib/python3.7/dist-packages (from google-auth<2,>=1.6.3->tensorboard==2.5) (0.2.8)\n",
            "Requirement already satisfied: rsa<5,>=3.1.4 in /usr/local/lib/python3.7/dist-packages (from google-auth<2,>=1.6.3->tensorboard==2.5) (4.7.2)\n",
            "Requirement already satisfied: requests-oauthlib>=0.7.0 in /usr/local/lib/python3.7/dist-packages (from google-auth-oauthlib<0.5,>=0.4.1->tensorboard==2.5) (1.3.0)\n",
            "Requirement already satisfied: importlib-metadata in /usr/local/lib/python3.7/dist-packages (from markdown>=2.6.8->tensorboard==2.5) (4.6.1)\n",
            "Requirement already satisfied: pyasn1<0.5.0,>=0.4.6 in /usr/local/lib/python3.7/dist-packages (from pyasn1-modules>=0.2.1->google-auth<2,>=1.6.3->tensorboard==2.5) (0.4.8)\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,>=2.21.0->tensorboard==2.5) (1.24.3)\n",
            "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests<3,>=2.21.0->tensorboard==2.5) (3.0.4)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests<3,>=2.21.0->tensorboard==2.5) (2.10)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests<3,>=2.21.0->tensorboard==2.5) (2021.5.30)\n",
            "Requirement already satisfied: oauthlib>=3.0.0 in /usr/local/lib/python3.7/dist-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<0.5,>=0.4.1->tensorboard==2.5) (3.1.1)\n",
            "Requirement already satisfied: typing-extensions>=3.6.4 in /usr/local/lib/python3.7/dist-packages (from importlib-metadata->markdown>=2.6.8->tensorboard==2.5) (3.7.4.3)\n",
            "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.7/dist-packages (from importlib-metadata->markdown>=2.6.8->tensorboard==2.5) (3.5.0)\n",
            "Installing collected packages: tensorboard\n",
            "  Attempting uninstall: tensorboard\n",
            "    Found existing installation: tensorboard 2.4.1\n",
            "    Uninstalling tensorboard-2.4.1:\n",
            "      Successfully uninstalled tensorboard-2.4.1\n",
            "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
            "pytorch-lightning 1.3.8 requires tensorboard!=2.5.0,>=2.2.0, but you have tensorboard 2.5.0 which is incompatible.\u001b[0m\n",
            "Successfully installed tensorboard-2.5.0\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "H_JmIQoiz6aG",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "3082fa9b-2ad8-42a9-832a-521cb06882d7"
      },
      "source": [
        "!pip install people_segmentation  > /dev/null"
      ],
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
            "tensorflow 2.5.0 requires tensorboard~=2.5, but you have tensorboard 2.4.1 which is incompatible.\u001b[0m\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bdS-B03Az-r-"
      },
      "source": [
        "from people_segmentation.pre_trained_models import create_model"
      ],
      "execution_count": 13,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7oGB-uNz0TWQ",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "1127fcfd-0588-4975-8a2d-df2a4dcb857d"
      },
      "source": [
        "model = create_model(\"Unet_2020-07-20\")"
      ],
      "execution_count": 14,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.7/dist-packages/torch/hub.py:480: UserWarning: Falling back to the old format < 1.6. This support will be deprecated in favor of default zipfile format introduced in 1.6. Please redo torch.save() to save it in the new zipfile format.\n",
            "  warnings.warn('Falling back to the old format < 1.6. This support will be '\n"
          ],
          "name": "stderr"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "S5qFkOsH_tAs"
      },
      "source": [
        "model.eval();"
      ],
      "execution_count": 15,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "W4qSnDthmpD4"
      },
      "source": [
        "import os\n",
        "import cv2\n",
        "from PIL import Image\n",
        "from google.colab import files \n",
        "import tarfile\n",
        "\n",
        "def make_targz_one_by_one(output_filename, source_dir):\n",
        "  tar = tarfile.open(output_filename,\"w\")\n",
        "  for root,dir_name,files_list in os.walk(source_dir):\n",
        "    for file in files_list:\n",
        "      pathfile = os.path.join(root, file)\n",
        "      tar.add(pathfile)\n",
        "  tar.close()\n",
        "  files.download(output_filename)"
      ],
      "execution_count": 16,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uVC7JJyw0VSe",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 242
        },
        "outputId": "b60956d8-660a-46ef-da93-f2f7b15e7246"
      },
      "source": [
        "path = 'ds7/'\n",
        "mask_path = 'mask7/'\n",
        "path_list = os.listdir(path)\n",
        "\n",
        "for i in path_list:\n",
        "  print(i)\n",
        "  image = load_rgb(path+i)\n",
        "  transform = albu.Compose([albu.Normalize(p=1)], p=1)\n",
        "  padded_image, pads = pad(image, factor=32, border=cv2.BORDER_CONSTANT)\n",
        "  x = transform(image=padded_image)[\"image\"]\n",
        "  x = torch.unsqueeze(tensor_from_rgb_image(x), 0)\n",
        "  with torch.no_grad():\n",
        "    prediction = model(x)[0][0]\n",
        "  mask = (prediction > 0).cpu().numpy().astype(np.uint8)\n",
        "  mask = unpad(mask, pads)\n",
        "  cv2.imwrite(mask_path+i, mask * 255)\n",
        "  # dst = cv2.addWeighted(image, 1, (cv2.cvtColor(mask, cv2.COLOR_GRAY2RGB) * (0, 255, 0)).astype(np.uint8), 0.5, 0)\n",
        "  # img = Image.fromarray(np.uint8(dst))\n",
        "  # img.save(result_path+i)\n",
        "\n",
        "make_targz_one_by_one('maskds7', 'mask7')"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "error",
          "ename": "NameError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-1-08e89d18a553>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0mpath\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m'ds7/'\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      2\u001b[0m \u001b[0mmask_path\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m'mask7/'\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 3\u001b[0;31m \u001b[0mpath_list\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mos\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlistdir\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpath\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      4\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mpath_list\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mNameError\u001b[0m: name 'os' is not defined"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kFj9ebO4remX"
      },
      "source": [
        "!rm -rf ds8/\n",
        "!rm -rf mask8/"
      ],
      "execution_count": 40,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "TzrM7xIBR25S",
        "outputId": "64a39e55-f554-437f-f44e-5b0f5feeba23"
      },
      "source": [
        "import glob\n",
        "path_file_number = glob.glob('ds7/*.jpeg') # 指定文件夹下文件个数\n",
        "print(len(path_file_number))"
      ],
      "execution_count": 41,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "232\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7JNTIcl80tUs"
      },
      "source": [
        "!mkdir /content/ds8\n",
        "!mkdir /content/mask8"
      ],
      "execution_count": 32,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_DqSY8vONEDI"
      },
      "source": [
        "!pip clone https://github.com/supervisely/supervisely"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2sTH6MFx1BRL"
      },
      "source": [
        "# 生成mask(label)\n",
        "\n",
        "import os\n",
        "import numpy as np\n",
        "import cv2\n",
        "import tqdm\n",
        "import supervisely_lib as sly\n",
        "from matplotlib import pyplot as plt\n",
        " \n",
        "def genMask(dataset_root_dir=\"/ds7\",\n",
        "            output_dir=\"/mask\",\n",
        "            label_color=[255,255,255]):\n",
        "    ''' 生成Mask\n",
        "        dataset_root_dir: 数据集根目录\n",
        "        output_dir: 生成mask的保存目录\n",
        "        label_color: mask 标注的颜色[R, G, B]\n",
        "    '''\n",
        "    project = sly.Project(dataset_root_dir, sly.OpenMode.READ)\n",
        " \n",
        "    # 打印数据集信息\n",
        "    print(f\"Project name     : {project.name}\")\n",
        "    print(f\"Project directory: {project.directory}\")\n",
        "    print(f\"Total images     : {project.total_items}\")\n",
        "    print(f\"Dataset names    : {project.datasets.keys()}\")\n",
        "    print()\n",
        "    print(project.meta)\n",
        " \n",
        "    pbar = tqdm.tqdm(total=project.total_items)\n",
        "    for dataset in project:\n",
        "        for item_name in dataset:\n",
        "            # 更新进度条\n",
        "            pbar.update(1)\n",
        "            # 获取原始图像和标注文件路径\n",
        "            item_paths = dataset.get_item_paths(item_name)\n",
        "            # 加载注释文件\n",
        "            ann = sly.Annotation.load_json_file(item_paths.ann_path, project.meta)\n",
        "            # 创建一个用于渲染标注的3通道黑色画布，\n",
        "            ann_render = np.zeros(ann.img_size + (3,), dtype=np.uint8)\n",
        "            # 渲染所有的标注（该数据集只有人）\n",
        "            ann.draw(ann_render, color=label_color)\n",
        "            # ann_render shape: (h, w, c), pixel: (R, G, B)\n",
        "            # RGB -> BGR，用于Opencv\n",
        "            ann_render = ann_render[..., ::-1]\n",
        "            # mask 保存目录不存在，则创建\n",
        "            save_dir = os.path.join(output_dir, dataset.name)\n",
        "            if not os.path.exists(save_dir):\n",
        "                os.makedirs(save_dir)\n",
        "            # 保存mask\n",
        "            mask_path = os.path.join(save_dir, item_name)\n",
        "            cv2.imwrite(mask_path, ann_render)\n",
        "    pbar.close()\n",
        "\n",
        "if __name__ == '__main__':\n",
        "    import fire\n",
        "    fire.Fire(genMask, name='gen_mask')"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}