fbrs_interactive_segmentation

f-BRS: 交互式分割的反向传播优化方法的重新思考 [论文] [PyTorch实现] [MXNet实现] [视频]

本代码库提供了训练和测试最先进的交互式分割模型的代码，是以下论文的官方PyTorch实现：

f-BRS: 交互式分割的反向传播优化方法的重新思考 Konstantin Sofiiuk、Ilia Petrov、Olga Barinova、Anton Konushin 三星研究院 https://arxiv.org/abs/2001.10331

请观看此视频了解我们算法的工作原理：

我们还提供了算法的完整MXNet实现，你可以查看mxnet分支。

新闻

[2021-02-16] 我们发布了一篇关于交互式分割的新论文（附带代码）：利用掩码引导恢复迭代训练以进行交互式分割。这是一种更简单的方法，达到了新的最佳性能，且无需任何测试时优化技术。

环境配置

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

pip3 install -r requirements.txt

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

如果你想进行训练或测试，必须在config.yml中配置数据集的路径（SBD用于训练和测试，GrabCut、Berkeley、DAVIS和COCO_MVal仅用于测试）。

交互式分割演示

图形界面基于TkInter库及其Python绑定。你可以使用任何提供的模型（见下文预训练模型部分）尝试交互式演示。我们的脚本会自动检测加载模型的架构，只需指定相应检查点的路径即可。

脚本使用示例：

# 此命令在GPU 0上运行交互式演示，使用cfg.INTERACTIVE_MODELS_PATH中的ResNet-34模型
# --checkpoint可以是相对于cfg.INTERACTIVE_MODELS_PATH的路径，也可以是检查点的绝对路径
python3 demo.py --checkpoint=resnet34_dh128_sbd --gpu=0

# 此命令使用/home/demo/fBRS/weights/中的ResNet-34模型运行交互式演示
# 如果你的GPU内存不足，可以减小--limit-longest-size（默认为800）
python3 demo.py --checkpoint=/home/demo/fBRS/weights/resnet34_dh128_sbd --limit-longest-size=400

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

你还可以使用Docker镜像运行演示。为此，你需要激活X-host连接，然后使用一些附加标志运行容器：

# 激活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

控制方式：

按鼠标左键和右键分别进行正面和负面点击；
滚轮可放大和缩小；
按住鼠标右键并拖动可在图像中移动（也可以使用方向键和WASD）；
按空格键完成当前对象；
打开多个文件时，按左箭头键显示上一张图像，按右箭头键显示下一张图像；
使用Ctrl+S保存当前正在编辑的注释（"原文件名".png）。

交互式分割选项：

放大功能（可通过复选框开启/关闭）
- 跳过点击次数 - 使用放大功能前需要跳过的点击次数。
- 目标尺寸 - 放大裁剪区域会被调整大小，使其较长边与此值匹配（对于大物体可增加此值）。
- 扩展比例 - 物体边界框在裁剪前会按此比例重新缩放。
BRS参数（可通过下拉菜单更改BRS类型）
- 网络点击次数 - 包含在网络输入中的前几次点击次数。后续的点击仅使用BRS处理（NoBRS忽略此选项）。
- L-BFGS-B最大迭代次数 - BRS中每步优化的最大函数评估次数（增加此值可提高准确度，但每次点击的计算时间会更长）。
可视化参数
- 预测阈值 滑块用于调整当前物体概率图的二值化阈值。
- Alpha混合系数 滑块用于调整所有预测掩码的强度。
- 可视化点击半径 滑块用于调整表示点击的红色和绿色圆点大小。

数据集

我们在SBD数据集上训练所有模型，并在GrabCut、Berkeley、DAVIS、SBD和COCO_MVal数据集上评估它们。我们还提供了在COCO和LVIS数据集组合上训练的模型结果。

Berkeley数据集包含100个实例（96张独特图像），由[K. McGuinness, 2010]提供。我们使用DAVIS数据集中相同的345张图像进行测试，与[WD Jang, 2019]一致，每张图像的真实标注掩码是所有物体掩码的并集。对于SBD数据集的测试，我们遵循[WD Jang, 2019]的协议，分别对测试集中的每个实例评估我们的算法。

为构建COCO_MVal数据集，我们从COCO 2017验证集中采样800个物体实例。具体来说，我们从80个类别中各采样10个独特实例。唯一的例外是烤面包机类别，在instances_val2017注释中只有9个独特实例。因此，为了得到800个掩码，其中一个类别包含11个物体。我们提供此数据集供下载，以便每个人都能复现我们的结果。

数据集	描述	下载链接
SBD	训练集：8498张图像，20172个实例<br>测试集：2857张图像，6671个实例	官方网站
Grab Cut	50张图像，每张一个物体	GrabCut.zip (11 MB)
Berkeley	96张图像，100个实例	Berkeley.zip (7 MB)
DAVIS	345张图像，每张一个物体	DAVIS.zip (43 MB)
COCO_MVal	800张图像，800个实例	COCO_MVal.zip (127 MB)

下载并解压后，别忘了在config.yml中更改数据集路径。

测试

预训练模型

我们提供了不同骨干网络的交互式分割预训练模型。评估结果与我们论文中的结果有所不同，因为我们在本仓库中展示的新代码基础上重新训练了所有模型。我们大大加速了RGB-BRS算法的推理速度 - 现在在SBD数据集上的运行速度比论文中给出的时间快2.5到4倍。尽管如此，新的结果有时甚至更好。

请注意，所有ResNet模型都是使用MXNet分支训练的，然后转换为PyTorch（它们的结果等效）。我们提供了用于转换模型的脚本。HRNet模型是使用PyTorch训练的。

你可以在下表中找到模型权重和测试结果：

<table> <tr> <th>骨干网络</th> <th>训练数据集</th> <th>链接</th> </tr> <tr> <td>ResNet-34</td> <td>SBD</td> <td><a href="https://github.com/saic-vul/fbrs_interactive_segmentation/releases/download/v1.0/resnet34_dh128_sbd.pth">resnet34_dh128_sbd.pth (GitHub, 89 MB)</a></td> </tr> <tr> <td>ResNet-50</td> <td>SBD</td> <td><a href="https://github.com/saic-vul/fbrs_interactive_segmentation/releases/download/v1.0/resnet50_dh128_sbd.pth">resnet50_dh128_sbd.pth (GitHub, 120 MB)</a></td> </tr> <tr> <td>ResNet-101</td> <td>SBD</td> <td><a href="https://github.com/saic-vul/fbrs_interactive_segmentation/releases/download/v1.0/resnet101_dh256_sbd.pth">resnet101_dh256_sbd.pth (GitHub, 223 MB)</a></td> </tr> <tr> <td>HRNetV2-W18+OCR</td> <td>SBD</td> <td><a href="https://github.com/saic-vul/fbrs_interactive_segmentation/releases/download/v1.0/hrnet18_ocr64_sbd.pth">hrnet18_ocr64_sbd.pth (GitHub, 39 MB)</a></td> </tr> <tr> <td>HRNetV2-W32+OCR</td> <td>SBD</td> <td><a href="https://github.com/saic-vul/fbrs_interactive_segmentation/releases/download/v1.0/hrnet32_ocr128_sbd.pth">hrnet32_ocr128_sbd.pth (GitHub, 119 MB)</a></td> </tr> <tr> <td>ResNet-50</td> <td>COCO+LVIS</td> <td><a href="https://github.com/saic-vul/fbrs_interactive_segmentation/releases/download/v1.0/resnet50_dh128_lvis.pth">resnet50_dh128_lvis.pth (GitHub, 120 MB)</a></td> </tr> <tr> <td>HRNetV2-W32+OCR</td> <td>COCO+LVIS</td> <td><a href="https://github.com/saic-vul/fbrs_interactive_segmentation/releases/download/v1.0/hrnet32_ocr128_lvis.pth">hrnet32_ocr128_lvis.pth (GitHub, 119 MB)</a></td> </tr> </table> <table align="center"> <tr> <th rowspan="2">模型</th> <th rowspan="2">BRS 类型</th> <th colspan="2">GrabCut</th> <th colspan="2">Berkeley</th> <th colspan="2">DAVIS</th> <th colspan="2">SBD</th> <th colspan="2">COCO_MVal</th> </tr> <tr> <td>NoC 85%</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> <td>NoC 85%</td> <td>NoC 90%</td> </tr> <tr> <td rowspan="2">ResNet-34 (SBD)</td> <td>RGB-BRS</td> <td>2.04</td> <td>2.50</td> <td>2.22</td> <td>4.49</td> <td>5.34</td> <td>7.91</td> <td>4.19</td> <td>6.83</td> <td>4.16</td> <td>5.52</td> </tr> <tr> <td>f-BRS-B</td> <td>2.06</td> <td>2.48</td> <td>2.40</td> <td>4.17</td> <td>5.34</td> <td>7.73</td> <td>4.47</td> <td>7.28</td> <td>4.31</td> <td>5.79</td> </tr> <tr> <td rowspan="2">ResNet-50 (SBD)</td> <td>RGB-BRS</td> <td>2.16</td> <td>2.56</td> <td>2.17</td> <td>4.27</td> <td>5.27</td> <td>7.51</td> <td>4.00</td> <td>6.59</td> <td>4.12</td> <td>5.61</td> </tr> <tr> <td>f-BRS-B</td> <td>2.20</td> <td>2.64</td> <td>2.17</td> <td>4.22</td> <td>5.44</td> <td>7.81</td> <td>4.55</td> <td>7.45</td> <td>4.31</td> <td>6.26</td> </tr> <tr> <td rowspan="2">ResNet-101 (SBD)</td> <td>RGB-BRS</td> <td>2.10</td> <td>2.46</td> <td>2.34</td> <td>3.91</td> <td>5.19</td> <td>7.23</td> <td>3.78</td> <td>6.28</td> <td>3.98</td> <td>5.45</td> </tr> <tr> <td>f-BRS-B</td> <td>2.30</td> <td>2.68</td> <td>2.61</td> <td>4.22</td> <td>5.32</td> <td>7.35</td> <td>4.20</td> <td>7.10</td> <td>4.11</td> <td>5.91</td> </tr> <tr> <td rowspan="2">HRNet-W18+OCR (SBD)</td> <td>RGB-BRS</td> <td>1.68</td> <td>1.94</td> <td>1.99</td> <td>3.81</td> <td>5.49</td> <td>7.98</td> <td>4.19</td> <td>6.84</td> <td>3.62</td> <td>5.04</td> </tr> <tr> <td>f-BRS-B</td> <td>1.86</td> <td>2.18</td> <td>2.07</td> <td>3.96</td> <td>5.62</td> <td>8.08</td> <td>4.70</td> <td>7.65</td> <td>3.87</td> <td>5.57</td> </tr> <tr> <td rowspan="2">HRNet-W32+OCR (SBD)</td> <td>RGB-BRS</td> <td>1.80</td> <td>2.16</td> <td>2.00</td> <td>3.58</td> <td>5.40</td> <td>7.59</td> <td>3.87</td> <td>6.33</td> <td>3.61</td> <td>5.12</td> </tr> <tr> <td>f-BRS-B</td> <td>1.78</td> <td>2.16</td> <td>2.13</td> <td>3.69</td> <td>5.54</td> <td>7.62</td> <td>4.31</td> <td>7.08</td> <td>3.82</td> <td>5.44</td> </tr> <tr> <td class="divider" colspan="12"><hr /></td> </tr> <tr> <td rowspan="2">ResNet-50 (COCO+LVIS)</td> <td>RGB-BRS</td> <td>1.54</td> <td>1.76</td> <td>1.56</td> <td>2.70</td> <td>4.93</td> <td>6.22</td> <td>4.04</td> <td>6.85</td> <td>2.41</td> <td>3.47</td> </tr> <tr> <td>f-BRS-B</td> <td>1.52</td> <td>1.74</td> <td>1.56</td> <td>2.61</td> <td>4.94</td> <td>6.36</td> <td>4.29</td> <td>7.20</td> <td>2.34</td> <td>3.43</td> </tr> <tr> <td rowspan="2">HRNet-W32+OCR (COCO+LVIS)</td> <td>RGB-BRS</td> <td>1.54</td> <td>1.60</td> <td>1.63</td> <td>2.59</td> <td>5.06</td> <td>6.34</td> <td>4.18</td> <td>6.96</td> <td>2.38</td> <td>3.55</td> </tr> <tr> <td>f-BRS-B</td> <td>1.54</td> <td>1.69</td> <td>1.64</td> <td>2.44</td> <td>5.17</td> <td>6.50</td> <td>4.37</td> <td>7.26</td> <td>2.35</td> <td>3.44</td> </tr> </table>

评估

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

以下命令运行模型评估（使用'-h'显示其他选项）：

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

脚本使用示例：

# 此命令在f-BRS-B模式下评估ResNet-34模型在所有数据集上的表现。
python3 scripts/evaluate_model.py f-BRS-B --checkpoint=resnet34_dh128_sbd

# 此命令在f-BRS-B模式下评估HRNetV2-W32+OCR模型在所有数据集上的表现。
python3 scripts/evaluate_model.py f-BRS-B --checkpoint=hrnet32_ocr128_sbd

# 此命令在RGB-BRS模式下评估ResNet-50模型在GrabCut和Berkeley数据集上的表现。
python3 scripts/evaluate_model.py RGB-BRS --checkpoint=resnet50_dh128_sbd --datasets=GrabCut,Berkeley

# 此命令在DistMap-BRS模式下评估ResNet-101模型在DAVIS数据集上的表现。
python3 scripts/evaluate_model.py DistMap-BRS --checkpoint=resnet101_dh256_sbd --datasets=DAVIS

Jupyter notebook

你还可以使用test_any_model.ipynb Jupyter notebook交互式地试验我们的模型。

训练

我们提供了在SBD数据集上训练我们模型的脚本。你可以使用以下命令开始训练：

# ResNet-34模型
python3 train.py models/sbd/r34_dh128.py --gpus=0,1 --workers=4 --exp-name=first-try

# ResNet-50模型
python3 train.py models/sbd/r50_dh128.py --gpus=0,1 --workers=4 --exp-name=first-try

# ResNet-101模型
python3 train.py models/sbd/r101_dh256.py --gpus=0,1,2,3 --workers=6 --exp-name=first-try

# HRNetV2-W32+OCR模型
python3 train.py models/sbd/hrnet32_ocr128.py --gpus=0,1 --workers=4 --exp-name=first-try

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

请注意，我们在2个GPU上训练了ResNet-34和ResNet-50模型，在4个GPU上训练了ResNet-101模型（我们使用Nvidia Tesla P40进行训练）。如果你打算在不同的GPU配置下训练模型，你需要设置不同的批量大小。你可以使用命令行参数--batch-size指定批量大小，或在模型脚本中更改默认值。

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

许可证

代码以MPL 2.0许可证发布。MPL是一个易于遵守的copyleft许可证。你必须根据MPL提供你所做任何更改的源代码，但你可以将MPL软件与专有代码结合，只要你将MPL代码保存在单独的文件中。

引用

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

@inproceedings{fbrs2020,
   title={f-brs: 重新思考用于交互式分割的反向传播细化},
   author={索菲乌克，康斯坦丁 和 彼得罗夫，伊利亚 和 巴里诺娃，奥尔加 和 科努申，安东},
   booktitle={IEEE/CVF计算机视觉与模式识别会议论文集},
   pages={8623--8632},
   year={2020}
}