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个物体。我们提供此数据集供下载，以便每个人都能复现我们的结果。

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

测试

预训练模型

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

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

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

评估

我们提供了脚本来测试所有呈现的模型在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}
}