UniverSeg:通用医学图像分割模型来了!

自从今年以来ChatGPT爆火和GPT-4的发布,一时间在大模型的潮流下,通用人工智能(AGI)也呼之欲出。随着本月初SAM和SegGPT等通用的CV大模型的提出,大模型和通用模型这把火也逐渐烧到的CV领域,特别是图像分割领域。很多做分割方向的小伙伴自我调侃说一觉醒来,自己的方向没了。

笔者所在的医学影像行业,一直以数据和高成本标注而筑起非常高的领域壁垒。几个月前要是有人跟我说想做一个医学影像的通用分割模型,我一定会觉得你是在说胡话。但此一时彼一时也,月初SAM发布的时候,主要对标的自然图像,笔者也测试了其在医学影像上表现,效果可以说超出预期了,但远不如自然图像,这让笔者感到医学影像的领域壁垒要打破可能还早。

但打脸来得太快。昨天看到了MIT发布的UniverSeg工作,名字很直接,针对的就是医学影像的通用分割:UniverSeg: Universal Medical Image Segmentation。趁着今天周末在家,赶紧把论文打印出来认真读了一遍。 

UniverSeg:通用医学图像分割模型来了!_第1张图片

近几年深度学习图像分割一直是高度定制化的,而医学影像分割尤其如此,对于不同影像模态(CT、MR、超声、X-ray、OCT等)、不同身体部位、不同标注(同一影像,可能会有不同的目标对象)。笔者所在的行业和方向,做一个目标影像或者疾病的分割任务,一定是专门收集该任务的影像数据,然后找专业人员来标注,再训练一个定制化的分割模型,这个模型是高度垂直的、定制化的和难以迁移的。但现在SAM和UnvierSeg的出现打破了这一范式。如图2所示,区别于此前的训练-预测的分割范式,UniverSeg给出的分割范式是:输入查询图像(Query Image),即待分割图像,再给定若干对提示图像-标注对(Support Set),UniverSeg就可以对查询图像进行很好的分割。 

UniverSeg:通用医学图像分割模型来了!_第2张图片

数据

先来看数据部分。UniverSeg研究总共使用了53个开源或者半开源的医学影像数据集,总共包括26个医学领域、16种影像模态和22000多次扫描。因为聚焦的是2D医学图像分割,所以对来源数据集中属于3D volume的数据都做了层采样,包括中间层采样(mid-slices)和最多标签体素层采样(max-slices)。基于这53个数据集构建的新数据被命名为MegaMedical数据集。但因为各来源数据集的有限获取规则,所以作者也没有开源MegaMedical。训练集构成信息如图3所示。

UniverSeg:通用医学图像分割模型来了!_第3张图片

模型

模型部分是UniverSeg工作的重点部分。传统的分割任务范式针对数据集,目标是训练一个监督学习模型,其中即为常规的卷积分割网络。但UniverSeg给出的模型范式是:,其中表示给定任务的输入,即待分割的查询图像,表示为提示图像和标注对:。

UniverSeg针对以上分割任务范式,提出了一种交叉卷积构成的CrossBlock结构,能够在查询图像和提示图像之间进行信息交互。给定查询图像特征图和一组提示图像特征图,交叉卷积层可以定义为:

其中表示两个特征图的concatenation。基于上述交叉层,进一步地可搭建交叉模块为:

最终,基于交叉卷积模块构建出的UniverSeg网络结构如图3所示。同时,为了提高UniverSeg的性能,模型在训练的时候也都做了数据增强,包括In-Task和Task两种数据增强方案。UniverSeg的训练流程如下图所示。整体结构仍然是基于UNet的编解码架构,只是卷积层都用CrossConv模块来代替了。

UniverSeg:通用医学图像分割模型来了!_第4张图片

实验

实验部分主要是与一些少样本模型(SENet、ALPNet、PANet)做了基线对比,然后在任务多样性、提示数据规模、查询图像规模等方面做了消融研究。UniverSeg试验部分做得很扎实,附录里面给了很多维度的性能结果,笔者这里不做一一展示。

UniverSeg:通用医学图像分割模型来了!_第5张图片

UniverSeg:通用医学图像分割模型来了!_第6张图片

总结

UniverSeg的主要贡献包括:

  • 提出了基于CrossBlock交叉模块机制的UniverSeg作为新的医学图像分割范式,无需重新训练新分割模型。

  • 与以往的少样本模型相比,UniverSeg在未经训练的新影像上表现出了与监督分割模型相当的精度。

  • 训练时的样本多样性能够使得UniverSeg有着强大的泛化性能。

总体来讲,在数据壁垒极高和标注成本极大的医学图像领域,通用分割模型不在遥不可及。随着AGI和CV大模型的发展,UniverSeg以及后续研究会逐渐推进这一领域的发展。2D分割模型也会逐渐向2.5D和3D扩展和延伸。

代码

项目代码地址:https://github.com/JJGO/UniverSeg demo地址:https://colab.research.google.com/drive/19Sauvhyzae5qvVLguaZRCuH1vJ5oTuw-?usp=sharing


另外,为了聚集更多的人参与到AI生产力工具上来,笔者前几天特意组建了一个名为【ChatGPT实验室】的知识星球,目前已有170+读者加入,星球的主要定位包括:

1. 如何基于ChatGPT提高工作和学习效率。

2. 跟踪NLP、LLM、AIGC和AGI的前沿动态和最新进展。

3. 分享ChatGPT的最新应用和玩法。

UniverSeg:通用医学图像分割模型来了!_第7张图片

你可能感兴趣的:(人工智能,机器学习,计算机视觉,深度学习)