通过掩码引导恢复迭代训练实现交互式分割

该仓库提供了训练和测试最先进的基于点击的交互式分割模型的源代码，这是以下论文的官方PyTorch实现：

通过掩码引导恢复迭代训练实现交互式分割 Konstantin Sofiiuk、Ilia Petrov、Anton Konushin 三星研究院 https://arxiv.org/abs/2102.06583

摘要： 最近关于基于点击的交互式分割的研究通过使用各种推理时优化方案展示了最先进的结果。这些方法相比前馈方法计算成本要高得多，因为它们在推理过程中需要通过网络执行反向传播，而且很难部署在通常只支持前向传播的移动框架上。在本文中，我们广泛评估了交互式分割的各种设计选择，发现无需任何额外的优化方案就可以获得新的最先进结果。因此，我们提出了一个简单的前馈模型用于基于点击的交互式分割，该模型利用了前几步的分割掩码。它不仅可以分割全新的对象，还可以从外部掩码开始并进行修正。在分析在不同数据集上训练的模型性能时，我们观察到训练数据集的选择极大地影响了交互式分割的质量。我们发现，在具有多样化和高质量注释的COCO和LVIS组合数据集上训练的模型表现优于所有现有模型。

配置环境

该框架基于Python 3.6构建，并依赖于PyTorch 1.4.0+。以下命令安装所有必要的包：

pip3 install -r requirements.txt

你也可以使用我们的Dockerfile来构建一个配置好环境的容器。

如果你想运行训练或测试，必须在config.yml中配置数据集的路径。

交互式分割演示

该GUI基于TkInter库及其Python绑定。你可以使用任何提供的模型尝试我们的交互式演示。我们的脚本会自动检测加载模型的架构，只需指定相应检查点的路径即可。

脚本使用示例：

# 该命令在GPU 0上运行HRNet18 ITER-M模型的交互式演示，模型路径为cfg.INTERACTIVE_MODELS_PATH
# --checkpoint可以是相对于cfg.INTERACTIVE_MODELS_PATH的路径，也可以是检查点的绝对路径
python3 demo.py --checkpoint=hrnet18_cocolvis_itermask_3p --gpu=0

# 该命令运行来自/home/demo/isegm/weights/的HRNet18 ITER-M模型的交互式演示
# 如果你的GPU内存不足，可以减小--limit-longest-size（默认为800）
python3 demo.py --checkpoint=/home/demo/fBRS/weights/hrnet18_cocolvis_itermask_3p --limit-longest-size=400

# 你也可以在纯CPU模式下尝试演示
python3 demo.py --checkpoint=hrnet18_cocolvis_itermask_3p --cpu

<details> <summary>在 Docker 中运行演示</summary> <pre><code># 激活 xhost xhost + docker run -v "$PWD":/tmp/ \ -v /tmp/.X11-unix:/tmp/.X11-unix \ -e DISPLAY=$DISPLAY <已构建的 Docker 镜像的 ID 或标签> \ python3 demo.py --checkpoint resnet34_dh128_sbd --cpu </code></pre> </details>

控制：

按键	描述
<kbd>鼠标左键</kbd>	放置一个正点击
<kbd>鼠标右键</kbd>	放置一个负点击
<kbd>滚轮</kbd>	放大和缩小图像
<kbd>鼠标右键</kbd> + <br> <kbd>移动鼠标</kbd>	移动图像
<kbd>空格键</kbd>	完成当前对象的遮罩

<details> <summary>使用外部分割遮罩初始化 ITER-M 模型</summary> <img src="https://yellow-cdn.veclightyear.com/0a4dffa0/46dd11b3-8370-4420-bd7f-1fed5ce8b2d1.jpg" alt="drawing" width="80%"/>

根据我们的论文，ITER-M 模型将图像、编码的用户输入和上一步的遮罩作为输入。此外，用户可以在放置任何点击之前使用外部遮罩初始化模型，并使用相同的界面修正这个遮罩。事实证明，我们的模型能够成功处理这种情况并使修改遮罩成为可能。

要使用外部遮罩初始化任何 ITER-M 模型，请使用菜单栏中的"加载遮罩"按钮。

</details> <details> <summary>交互式分割选项</summary> <ul> <li>放大（可以使用复选框开启/关闭）</li> <ul> <li>跳过点击 - 在使用放大功能之前要跳过的点击次数。</li> <li>目标大小 - 放大裁剪后调整大小，使其较长边匹配此值（对于大物体可增加此值）。</li> <li>扩展比例 - 在裁剪前，对象边界框按此比例调整。</li> <li>固定裁剪 - 放大裁剪调整为（目标大小，目标大小）。</li> </ul> <li>BRS 参数（可以使用下拉菜单更改 BRS 类型）</li> <ul> <li>网络点击 - 包含在网络输入中的首次点击次数。后续点击仅使用 BRS 处理（NoBRS 忽略此选项）。</li> <li>L-BFGS-B 最大迭代次数 - BRS 中每一步优化的函数评估最大次数（增加可提高准确性，但每次点击的计算时间会更长）。</li> </ul> <li>可视化参数</li> <ul> <li>预测阈值 滑块调整当前对象概率图二值化的阈值。</li> <li>Alpha 混合系数 滑块调整所有预测遮罩的强度。</li> <li>可视化点击半径 滑块调整表示点击的红点和绿点的大小。</li> </ul> </ul> </details>

数据集

我们在 SBD 和 COCO+LVIS 上训练所有模型，并在 GrabCut、Berkeley、DAVIS、SBD 和 PascalVOC 上评估它们。我们还提供了额外数据集的链接：ADE20k 和 OpenImages，这些数据集用于消融研究。

数据集	描述	下载链接
ADE20k	22k 张图像，共 434k 个实例	[官方网站][ADE20k]
OpenImages	944k 张图像，共 2.6M 个实例	[官方网站][OpenImages]
MS COCO	118k 张图像，1.2M 个实例（训练集）	[官方网站][MSCOCO]
LVIS v1.0	100k 张图像，共 1.2M 个实例	[官方网站][LVIS]
COCO+LVIS*	99k 张图像，1.5M 个实例（训练集）	[原始 LVIS 图像][LVIS] + <br> [我们的组合注释][COCOLVIS_annotation]
SBD	8498 张图像，20172 个实例（训练集）<br>2857 张图像，6671 个实例（测试集）	[官方网站][SBD]
Grab Cut	50 张图像，每张一个对象（测试集）	[GrabCut.zip (11 MB)][GrabCut]
Berkeley	96 张图像，100 个实例（测试集）	[Berkeley.zip (7 MB)][Berkeley]
DAVIS	345 张图像，每张一个对象（测试集）	[DAVIS.zip (43 MB)][DAVIS]
Pascal VOC	1449 张图像，3417 个实例（验证集）	[官方网站][PascalVOC]
COCO_MVal	800 张图像，800 个实例（测试集）	[COCO_MVal.zip (127 MB)][COCO_MVal]
[ADE20k]: http://sceneparsing.csail.mit.edu/
[OpenImages]: https://storage.googleapis.com/openimages/web/download.html
[MSCOCO]: https://cocodataset.org/#download
[LVIS]: https://www.lvisdataset.org/dataset
[SBD]: http://home.bharathh.info/pubs/codes/SBD/download.html
[GrabCut]: https://github.com/saic-vul/fbrs_interactive_segmentation/releases/download/v1.0/GrabCut.zip
[Berkeley]: https://github.com/saic-vul/fbrs_interactive_segmentation/releases/download/v1.0/Berkeley.zip
[DAVIS]: https://github.com/saic-vul/fbrs_interactive_segmentation/releases/download/v1.0/DAVIS.zip
[PascalVOC]: http://host.robots.ox.ac.uk/pascal/VOC/
[COCOLVIS_annotation]: https://github.com/saic-vul/ritm_interactive_segmentation/releases/download/v1.0/cocolvis_annotation.tar.gz
[COCO_MVal]: https://github.com/saic-vul/fbrs_interactive_segmentation/releases/download/v1.0/COCO_MVal.zip

下载并解压数据集后，请记得在config.yml中更改数据集的路径。

(*) 要准备COCO+LVIS，您需要下载原始的LVIS v1.0，然后下载并解压我们预处理的注释文件，这些文件是通过将COCO和LVIS数据集合并到LVIS v1.0的文件夹中得到的。

测试

预训练模型

我们提供了具有不同骨干网络的交互式分割预训练模型。

您可以在下面的表格中找到模型权重和评估结果：

<table> <thead align="center"> <tr> <th rowspan="2">训练 数据集</th> <th rowspan="2">模型</th> <th colspan="2">GrabCut</th> <th>Berkeley</th> <th colspan="2">SBD</th> <th colspan="2">DAVIS</th> <th>Pascal VOC</th> <th>COCO MVal</th> </tr> <tr> <td>NoC 85%</td> <td>NoC 90%</td> <td>NoC 90%</td> <td>NoC 85%</td> <td>NoC 90%</td> <td>NoC 85%</td> <td>NoC 90%</td> <td>NoC 85%</td> <td>NoC 90%</td> </tr> </thead> <tbody align="center"> <tr> <td rowspan="1">SBD</td> <td align="left"><a href="https://github.com/saic-vul/ritm_interactive_segmentation/releases/download/v1.0/sbd_h18_itermask.pth">HRNet18 IT-M (38.8 MB)</a></td> <td>1.76</td> <td>2.04</td> <td>3.22</td> <td>3.39</td> <td>5.43</td> <td>4.94</td> <td>6.71</td> <td><ins>2.51</ins></td> <td>4.39</td> </tr> <tr> <td rowspan="4">COCO+ LVIS</td> <td align="left"><a href="https://github.com/saic-vul/ritm_interactive_segmentation/releases/download/v1.0/coco_lvis_h18_baseline.pth">HRNet18 (38.8 MB)</a></td> <td>1.54</td> <td>1.70</td> <td>2.48</td> <td>4.26</td> <td>6.86</td> <td>4.79</td> <td>6.00</td> <td>2.59</td> <td>3.58</td> </tr> <tr> <td align="left"><a href="https://github.com/saic-vul/ritm_interactive_segmentation/releases/download/v1.0/coco_lvis_h18s_itermask.pth">HRNet18s IT-M (16.5 MB)</a></td> <td>1.54</td> <td>1.68</td> <td>2.60</td> <td>4.04</td> <td>6.48</td> <td>4.70</td> <td>5.98</td> <td>2.57</td> <td>3.33</td> </tr> <tr> <td align="left"><a href="https://github.com/saic-vul/ritm_interactive_segmentation/releases/download/v1.0/coco_lvis_h18_itermask.pth">HRNet18 IT-M (38.8 MB)</a></td> <td>1.42</td> <td>1.54</td> <td><ins>2.26</ins></td> <td>3.80</td> <td>6.06</td> <td><ins>4.36</ins></td> <td><ins>5.74</ins></td> <td>2.28</td> <td><ins>2.98</ins></td> </tr> <tr> <td align="left"><a href="https://github.com/saic-vul/ritm_interactive_segmentation/releases/download/v1.0/coco_lvis_h32_itermask.pth">HRNet32 IT-M (119 MB)</a></td> <td><ins>1.46</ins></td> <td><ins>1.56</ins></td> <td>2.10</td> <td><ins>3.59</ins></td> <td><ins>5.71</ins></td> <td>4.11</td> <td>5.34</td> <td>2.57</td> <td>2.97</td> </tr> </tbody> </table>

评估

我们提供了脚本来测试所有呈现的模型在GrabCut、Berkeley、DAVIS、Pascal VOC和SBD上的所有可能配置。要测试模型，您应该下载其权重并将它们放在./weights文件夹中（您可以在config.yml中更改此路径，请参阅INTERACTIVE_MODELS_PATH变量）。要测试我们的任何模型，只需指定相应检查点的路径。我们的脚本会自动检测加载模型的架构。

以下命令在所有测试数据集上运行NoC评估（使用'-h'显示其他选项）：

python3 scripts/evaluate_model.py <brs-mode> --checkpoint=<checkpoint-name>

脚本使用示例：

# 此命令在NoBRS模式下评估HRNetV2-W18-C+OCR ITER-M模型在所有数据集上的性能。
python3 scripts/evaluate_model.py NoBRS --checkpoint=hrnet18_cocolvis_itermask_3p

# 此命令在f-BRS-B模式下评估HRNet-W18-C-Small-v2+OCR ITER-M模型在所有数据集上的性能。
python3 scripts/evaluate_model.py f-BRS-B --checkpoint=hrnet18s_cocolvis_itermask_3p

此命令在GrabCut和Berkeley数据集上以NoBRS模式评估HRNetV2-W18-C+OCR ITER-M模型。

python3 scripts/evaluate_model.py NoBRS --checkpoint=hrnet18_cocolvis_itermask_3p --datasets=GrabCut,Berkeley


### Jupyter笔记本

您还可以使用[test_any_model.ipynb](https://github.com/SamsungLabs/ritm_interactive_segmentation/blob/master/./notebooks/test_any_model.ipynb) Jupyter笔记本交互式地试验我们的模型。

## 训练

我们提供了在SBD数据集上训练我们模型的脚本。您可以使用以下命令开始训练：
```.bash
# ResNet-34非迭代基线模型
python3 train.py models/noniterative_baselines/r34_dh128_cocolvis.py --gpus=0 --workers=4 --exp-name=first-try

# HRNet-W18-C-Small-v2+OCR ITER-M模型
python3 train.py models/iter_mask/hrnet18s_cocolvis_itermask_3p.py --gpus=0 --workers=4 --exp-name=first-try

# HRNetV2-W18-C+OCR ITER-M模型
python3 train.py models/iter_mask/hrnet18_cocolvis_itermask_3p.py --gpus=0,1 --workers=6 --exp-name=first-try

# HRNetV2-W32-C+OCR ITER-M模型
python3 train.py models/iter_mask/hrnet32_cocolvis_itermask_3p.py --gpus=0,1,2,3 --workers=12 --exp-name=first-try

对于每个实验，都会在./experiments中创建一个单独的文件夹，包含Tensorboard日志、文本日志、可视化结果和检查点。您可以在config.yml中指定另一个路径（参见EXPS_PATH变量）。

请注意，我们在1个GPU上训练了ResNet-34和HRNet-18s，在2个GPU上训练了HRNet-18，在4个GPU上训练了HRNet-32（我们使用Nvidia Tesla P40进行训练）。如果要在不同的GPU上训练，您应该使用命令行参数--batch-size调整批量大小，或在模型脚本中更改默认值。

我们使用了来自官方仓库的预训练HRNetV2模型。如果您想使用这些模型进行交互式分割训练，需要下载权重并在config.yml中指定它们的路径。

许可证

该代码以MIT许可证发布。这是一个简短、宽松的软件许可证。基本上，只要您在软件/源代码的任何副本中包含原始版权和许可声明，您就可以随心所欲地使用它。

引用

如果您发现这项工作对您的研究有用，请引用我们的论文：

@inproceedings{ritm2022,
  title={Reviving iterative training with mask guidance for interactive segmentation},
  author={Sofiiuk, Konstantin and Petrov, Ilya A and Konushin, Anton},
  booktitle={2022 IEEE International Conference on Image Processing (ICIP)},
  pages={3141--3145},
  year={2022},
  organization={IEEE}
}

@inproceedings{fbrs2020,
   title={f-brs: Rethinking backpropagating refinement for interactive segmentation},
   author={Sofiiuk, Konstantin and Petrov, Ilia and Barinova, Olga and Konushin, Anton},
   booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
   pages={8623--8632},
   year={2020}
}