{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "69debbc2-30e8-4390-805a-cb6ebad35f3f",
   "metadata": {
    "tags": []
   },
   "source": [
    "## EasyCV图像分割-Mask2Former\n",
    "本文将介绍如何利用EasyCV使用Transformer-based图像分割算法[Mask2Former](https://arxiv.org/pdf/2112.01527.pdf)进行图像分割模型的训练,以及如何利用训练好的模型进行图像分割预测\n",
    " \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9ef9614-cb63-487f-bafb-c62c90ae607b",
   "metadata": {},
   "source": [
    "## 运行环境要求\n",
    "\n",
    "PAI-Pytorch镜像 or 原生Pytorch1.8+以上环境 GPU机器， 内存32G以上"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f855f736-c08a-44c2-b1c1-33a8e7043864",
   "metadata": {
    "tags": []
   },
   "source": [
    "## 安装依赖包\n",
    "\n",
    "注: 在PAI-DSW docker中无需安装相关依赖，可跳过此部分 在本地notebook环境中执行\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43e60ee5-f52f-4045-bc69-4a123731a5a5",
   "metadata": {},
   "source": [
    "1、 首先，安装pytorch和对应版本的torchvision，支持Pytorch1.8以上版本"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f97b2295-4d67-4a83-8637-ecde48fa3001",
   "metadata": {},
   "outputs": [],
   "source": [
    "# install pytorch and torch vision\n",
    "! conda install --yes pytorch==1.10.0 torchvision==0.11.0 -c pytorch"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff3ed9f8-703d-416e-8b70-10a8ce029e8d",
   "metadata": {},
   "source": [
    "2、获取torch和cuda版本，安装对应版本的mmcv和nvidia-dali"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c95571d5-3065-40b3-ad53-dc3063db8604",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import os\n",
    "os.environ['CUDA']='cu' + torch.version.cuda.replace('.', '')\n",
    "os.environ['Torch']='torch'+torch.version.__version__.replace('+PAI', '')\n",
    "!echo \"cuda version: $CUDA\"\n",
    "!echo \"pytorch version: $Torch\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce3dca16-9baf-42f8-94c4-5ea7f087d61e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# install some python deps\n",
    "! pip install mmdet\n",
    "! pip install mmcv-full==1.4.4 -f https://download.openmmlab.com/mmcv/dist/${CUDA}/${Torch}/index.html\n",
    "! pip install http://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/third_party/nvidia_dali_cuda100-0.25.0-1535750-py3-none-manylinux2014_x86_64.whl"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3dee3a99-191a-4515-97dd-80841db43775",
   "metadata": {},
   "source": [
    "3、  安装EasyCV算法包"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f03f9e9-6029-4918-b8dc-533db9c7fae3",
   "metadata": {},
   "outputs": [],
   "source": [
    "pip install pai-easycv"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3483904d-91dd-44c4-a486-246aa38d4124",
   "metadata": {},
   "source": [
    "4、 简单验证"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d382373-76a0-4fbc-b09f-5d082aab5104",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from easycv.apis import *"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6da90a06",
   "metadata": {},
   "source": [
    "5、安装deformable_attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a5d35b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import easycv\n",
    "print(easycv.__file__)\n",
    "# 进入easycv安装目录编译deformable_attention\n",
    "! cd /home/pai/lib/python3.6/site-packages/easycv/thirdparty/deformable_attention && python setup.py build install "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29bb3d55-d00b-453b-9522-c686260e325c",
   "metadata": {},
   "source": [
    "## 数据准备\n",
    "\n",
    "接下来介绍基于coco数据集的实例分割训练示例，你可以下载[COCO2017](https://cocodataset.org/#download)数据，也可以使用我们提供了示例COCO数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9492b324-17d9-4963-b7dd-a49f684c54c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "! wget http://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/data/small_coco_demo/small_coco_demo.tar.gz && tar -zxf small_coco_demo.tar.gz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5fdc9dc3-0886-493f-be5b-f858eff72164",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 重命名文件夹\n",
    "! mkdir -p data/  && mv small_coco_demo database/coco"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dfcdc227",
   "metadata": {},
   "source": [
    "data/coco格式如下\n",
    "\n",
    "```shell\n",
    "data/coco/\n",
    "├── annotations\n",
    "│   ├── instances_train2017.json\n",
    "│   └── instances_val2017.json\n",
    "├── train2017\n",
    "│   ├── 000000005802.jpg\n",
    "│   ├── 000000060623.jpg\n",
    "│   ├── 000000086408.jpg\n",
    "│   ├── 000000118113.jpg\n",
    "│   ├── 000000184613.jpg\n",
    "│   ├── 000000193271.jpg\n",
    "│   ├── 000000222564.jpg\n",
    "│       ...\n",
    "│   └── 000000574769.jpg\n",
    "└── val2017\n",
    "    ├── 000000006818.jpg\n",
    "    ├── 000000017627.jpg\n",
    "    ├── 000000037777.jpg\n",
    "    ├── 000000087038.jpg\n",
    "    ├── 000000174482.jpg\n",
    "    ├── 000000181666.jpg\n",
    "    ├── 000000184791.jpg\n",
    "    ├── 000000252219.jpg\n",
    "         ...\n",
    "    └── 000000522713.jpg\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9bd7101-3072-417c-ac54-5d7d254b31b4",
   "metadata": {},
   "source": [
    "## 模型训练\n",
    "\n",
    "这个Demo中我们采用[Mask2Former](https://arxiv.org/pdf/2112.01527.pdf)图像分割算法训练ResNet50主干网络， 下载示例配置文件"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab669385-8525-4694-8388-148ba1c2753a",
   "metadata": {},
   "outputs": [],
   "source": [
    "! rm -rf mask2former_r50_8xb2_e50_instance.py\n",
    "! wget https://raw.githubusercontent.com/alibaba/EasyCV/master/configs/segmentation/mask2former/mask2former_r50_8xb2_e50_instance.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2890d267-6b95-47e6-9b51-83402446fa7f",
   "metadata": {},
   "source": [
    "为了适配小数据，我们对配置文件mask2former_r50_8xb2_e50_instance.py做如下字段的修改，减少训练epoch数目，加大打印日志的频率\n",
    "\n",
    "```python\n",
    "\n",
    "total_epochs = 3\n",
    "\n",
    "#optimizer.lr -> 0.000001\n",
    "# optimizer\n",
    "optimizer = dict(\n",
    "    type='AdamW',\n",
    "    lr=0.000001,\n",
    "    weight_decay=0.05,\n",
    "    eps=1e-8,\n",
    "    betas=(0.9, 0.999),\n",
    "    paramwise_options={\n",
    "        'backbone': dict(lr_mult=0.1),\n",
    "        'query_embed': dict(weight_decay=0.),\n",
    "        'query_feat': dict(weight_decay=0.),\n",
    "        'level_embed': dict(weight_decay=0.),\n",
    "        'norm': dict(weight_decay=0.),\n",
    "    })\n",
    "\n",
    "# log_config.interval 1\n",
    "log_config = dict(interval=1)\n",
    "\n",
    "```\n",
    "\n",
    "注意： 如果是使用COCO完整数据训练，为了保证效果，建议使用单机8卡进行训练；\n",
    "\n",
    "为了保证模型效果，我们在[预训练模型](http://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/EasyCV/modelzoo/segmentation/mask2former_r50_instance/epoch_50.pth)基础上finetune， 执行如下命令启动训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aaf88f46-a578-4dfc-a33f-afa0357f734a",
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -m easycv.tools.train mask2former_r50_8xb2_e50_instance.py --work_dir work_dir/segmentatino/mask2former_r50_instance --load_from http://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/EasyCV/modelzoo/segmentation/mask2former_r50_instance/epoch_50.pth"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc43e194-d8e7-4796-af3a-1f64663b9744",
   "metadata": {},
   "source": [
    "### 预测"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2cc9e6fc",
   "metadata": {},
   "source": [
    "下载测试图片"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "973d5bd4",
   "metadata": {},
   "outputs": [],
   "source": [
    "! wget http://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/EasyCV/modelzoo/segmentation/000000123213.jpg"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ecb723f",
   "metadata": {},
   "source": [
    "使用训练好的模型进行图像分割预测"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a5a3632",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import cv2\n",
    "from IPython.display import Image\n",
    "from easycv.predictors.segmentation import Mask2formerPredictor\n",
    "\n",
    "\n",
    "predictor = Mask2formerPredictor(model_path='work_dir/segmentatino/mask2former_r50_instance/epoch_3.pth',\n",
    "                                 config_file='mask2former_r50_8xb2_e50_instance.py',\n",
    "                                 task_mode='instance')\n",
    "img = cv2.imread('000000123213.jpg')\n",
    "predict_out = predictor(['000000123213.jpg'])\n",
    "instance_img = predictor.show_instance(img, **predict_out[0])\n",
    "cv2.imwrite('instance_out.jpg',instance_img)\n",
    "display(Image('000000123213.jpg'))\n",
    "display(Image('instance_out.jpg'))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.6.13 ('torch1.10')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.13"
  },
  "vscode": {
   "interpreter": {
    "hash": "ffac244d5fb3e091416ac35ee470bc03f8b6d092e3cccc2d90f82acef7653459"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}