机器之心专栏
作者:Xuan Liao、Wenhao Li等
如何提高交互式图像分割算法的效率?上海交大和华师大的研究者提出了一种基于多智能体深度强化学习的新型算法。
现有的交互式图像分割算法虽然能迭代式地更新分割结果,但很大程度上忽略了对连续交互之间动态性的探索,造成分割效率大大降低。
在 CVPR 2020 的一篇论文中,来自上海交大和华师大的团队联合提出了一种基于多智能体深度强化学习(MARL)的新型交互式三维医疗图像分割算法(IteR-MRL)。通过将迭代更新的交互式图像分割的动态过程建模成马尔可夫过程,并使用 MARL 解决,IteR-MRL 实现了更少的交互次数和更快的收敛速度,在多个医疗图像数据集上超过了现有算法。
![4436dc65f8079d97dabf98bfa86b40a5.png](https://i-blog.csdnimg.cn/blog_migrate/be5f3b41577bf5c195f74fef9a430275.jpeg)
论文地址:https://arxiv.org/abs/1911.10334
现有交互式图像分割策略的缺点
目前的三维图像自动分割算法很难达到医用标准。为了得到更佳的分割结果,交互式的图像分割策略成为有价值的研究方向,此类策略通过引入少量的用户提示实现对分割结果的迭代优化。
现有的交互式算法虽然能迭代式地对分割结果进行多轮更新,但它们仍然独立地考虑每一轮更新的分割结果,很大程度上忽略了连续交互的动态性。
该研究的贡献
为了更好地利用交互式图像分割的动态性,来自上海交大和华师大的团队提出了一个基于深度强化学习的算法 IteR-MRL,将交互式医疗图像分割的动态过程建模成一个马尔科夫决策过程,然后用深度强化学习求解。该算法从整体上考虑分割更新序列,充分挖掘了交互分割前后的关联。
由于体素级的图像分割标注任务会造成单智能体的强化学习算法遭遇探索空间爆炸问题,因此该团队采用多智能体模型,将每个图像体素看作一个智能体。通过让所有智能体共享同一个行为策略,将探索空间减小到了可行的范围。多智能体模型还能同时实现图像体素的合作交互,从而有效地捕捉分割任务中体素之间的依赖性。
该研究的贡献主要有以下三点:
(1)将交互式图像分割任务建模成一个马尔科夫决策过程,提出了全新的基于 MARL 的交互式三维医疗图像分割框架,使得用户的交互得到更有效地利用;
(2)通过分割概率的形式来保留分割结果的不确定性,丰富了之前分割结果的信息,也实现了更加精确和精细的分割调整;
(3)实验表明,通过考虑连续两次分割结果之间的相对增益,分割效果得到显著提升,并具备更少的交互次数和更快的收敛速度。
方法概览
![104cc7757eef17dff099d5cc9a93c6fc.png](https://i-blog.csdnimg.cn/blog_migrate/9ccedffeec7889e64b2e8a8a26b7e47b.jpeg)
迭代式交互图像分割方法的流程图。
该研究采用迭代更新的流程。首先向模型提供初始的粗分割结果(可以来自任意图像分割算法或不精确的人工标注)。如果当前分割结果不符合预期,则用户在错误区域标注关键性的提示信息,更新算法根据用户提示再次更新分割结果。上述过程不断迭代,直至用户对优化后的分割结果满意为止。
基于 MARL 的交互式图像分割框架(IteR-MRL)
![b38354e06bb6fdd6a5568789f1bb5017.png](https://i-blog.csdnimg.cn/blog_migrate/33ad7dbb9169b9caf4e22bd1dd6b301c.jpeg)
IteR-MRL 算法图示。
在每一步交互更新中,将包含图像、之前分割结果和提示图的当前状态输入到 FCN 网络中,然后网络根据其输出的动作产生当前的分割结果。接着,用户根据错误区域给出的提示点(红点),进而生成新的提示图。每一步的回报值都由环境决定,即分割的预测标签和正确标签。在 MARL 的设定中,体素被看作是互相协作的智能体。
1. 状态设计
![2967cb73dc3f2521502676d7a634c69f.png](https://i-blog.csdnimg.cn/blog_migrate/d11e8b636c68c2ba1bb2612ed437ed10.jpeg)
每个智能体的状态由体素值、之前分割概率和提示图上对应的值组成,两张提示图分别代表了用户给出的正负提示。分割概率能够保留模型之前输出结果的不确定性。
2. 动作设计
![393ce5bb9e9b892fdc228c5ba7627efd.png](https://i-blog.csdnimg.cn/blog_migrate/c48e7d04e084b19e05ea300e6517631f.jpeg)
基于之前分割概率,每个智能体的动作是预测基于之前分割结果的改变量。将改变量施加到之前分割结果上,生成新的分割结果。动作空间由 K 个动作组成,使得智能体可以在不同情况下进行有针对性的多样化调整。
3. 回报设计
![c21d8e4c61384a73824a62c7f14dcb60.png](https://i-blog.csdnimg.cn/blog_migrate/e22c0fb0a7a02b36c64ee111989c9012.jpeg)
每个智能体的回报被设计为当前预测结果与正确标签之间的交叉熵变化量。这样的相对性能回报是基于性能变化趋势的,我们可以同时得知性能变化的方向及程度。比起一个遥远的目标,相对性能可以给智能体提供一条可供对比的基线,并超越。
实验
对比实验
1. 与不同初始分割算法的组合
![23aa9b6b799f15fce03576a2fa8faac2.png](https://i-blog.csdnimg.cn/blog_migrate/9d42d58d7940321613714a55123dbb86.jpeg)
更新算法以四种不同的初始分割算法的分割结果作为其初始分割结果。比起其他三种现有的更新算法,不论初始分割算法如何选择,IteR-MRL 都拥有更好的性能。这表明了该算法的鲁棒性和普适性。
2. 一个交互序列中的性能提升
![d85b209325cd657c221ae004d049f674.png](https://i-blog.csdnimg.cn/blog_migrate/331302e2632b46fdf4f1f692e32ed023.jpeg)
在其他算法的每步性能提升缓慢时,IteR-MRL 仍能保持相对高的性能提升,这证明了在交互式分割任务中考虑动态性可以有效且快速地提升性能。每一步相对高的性能提升逐渐累积,使得最终算法达到好的分割精度。
3. 分割结果可视化对比
![69adcfcad0883f3f8e85885b4c8a5d35.png](https://i-blog.csdnimg.cn/blog_migrate/4fd44a8f743c85fcc04fdf02335d5135.jpeg)
从上图中可以观察到,当其他算法倾向于产生相对光滑的分割轮廓时,IteR-MRL 算法在捕捉物体边缘细节信息时表现更佳。
4. 不同数据集的性能对比
![166d4f3c91a7e4807674debebd6fbf06.png](https://i-blog.csdnimg.cn/blog_migrate/b0c08e559b69e073cd09343fa5c19a25.jpeg)
该实验证明 IteR-MRL 算法对于数据集的鲁棒性,在脑部肿瘤、心脏和前列腺数据集上都拥有稳定的表现。
控制变量实验
1. 不同动作与状态设定的组合
![2ff580e6fc22b9f568ccd1f19c9d4946.png](https://i-blog.csdnimg.cn/blog_migrate/06cb54782df29678f00dbe9788b2127e.jpeg)
数值小的动作比数值大的动作拥有更好的性能,且能够提供多样化的调整尺度,使得模型动作更加充足,从而达到好的性能。
2. 回报的设计对性能的影响
![420a0e6cdb02c6e3bea36156814be55d.png](https://i-blog.csdnimg.cn/blog_migrate/0aa9de528c0d5e1ef62eb77453f79661.jpeg)
基于相对性能增益的回报具备更优秀的效果,原因是相对增益能够更好地反映智能体对分割概率的调整。