UniverSeg：通用医学图像分割模型来了！

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

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

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

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

数据

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

模型

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

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

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

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