[论文阅读] CarveMix: A Simple Data Augmentation Method for Brain Lesion Segmentation

[论文地址][代码][MICCAI 21]

Abstract

脑部病变分割(Brain Lesion Segmentation)为临床诊断提供了宝贵的工具，卷积神经网络(CNN)在该任务中取得了空前的成功。数据增强是一种广泛使用的策略，可以提高CNN的训练效果，而用于脑部病变分割的增强方法的设计仍然是一个开放的问题。在这项工作中，我们提出了一种简单的数据增强方法，被称为CarveMix，用于基于CNN的脑部病变分割。与其他基于"混合"的方法一样，如Mixup和CutMix，CarveMix随机地结合两个现有的标记图像来生成新的标记样本。然而，与这些基于图像组合的增强策略不同，CarveMix是病变感知的，在进行组合时要注意病变，并为生成的图像创建适当的标注。具体来说，我们根据病变的位置和几何形状，从一个标记的图像中刻出一个感兴趣的区域(ROI)，ROI的大小从一个概率分布中取样。然后，雕刻好的ROI取代了第二张标记图像中的相应体素，第二张图像的注释也相应被取代。通过这种方式，我们为网络训练生成新的标记图像，并且保留了病变信息。为了评估所提出的方法，我们在两个大脑病变数据集上进行了实验。结果表明，与其他简单的数据增强方法相比，我们的方法提高了分割的准确性。

Method

本文是一种专门为脑部病变分割提出的数据增强方法——CarveMix。该方法也是一种基于标签融合的方法，比如MixUp(ICLR 18)就是将两个标签给线性融合起来，而CarveMix(ICCV 19)则是一种非线性的融合。需要注意的是，这些经典的方法都是用于图像分类任务的，因此目前是比较缺乏用于分割任务的标签融合方法的。CarveMix的融合流程如下：

先直接看公式：融合得到的图像$\mathbf{X}$与融合得到的标签$\mathbf{Y}$的最终计算过程如下：$$\mathbf{X}={\mathbf{X}}_i\odot {\mathbf{M}}_i+{\mathbf{X}}_j\odot \left(1-{\mathbf{M}}_i\right)$$$$\mathbf{Y}={\mathbf{Y}}_i\odot {\mathbf{M}}_i+{\mathbf{Y}}_j\odot \left(1-{\mathbf{M}}_i\right)$$ 为了直观起见，我们以标签融合为例展示具体的融合过程，图像本身的融合与标签的融合在过程上是一致的。从图中看，标签的融合基本上就约等于直接把两张原始图像的mask ${\mathbf{Y}}_i$和${\mathbf{Y}}_j$直接给加起来：

而实际上也是${\mathbf{Y}}_i$乘了一个系数$a$后与${\mathbf{Y}}_j$乘了一个系数$b$后再相加，有$a+b=1$。图中"⊙"符号表示逐像素乘，因此甚至也可以将${\mathbf{M}}_i$视为一种空间注意力图。

现在的问题其实就是怎么计算${\mathbf{M}}_i$了。从图中可以看到，${\mathbf{M}}_i$与${\mathbf{Y}}_i$其实是非常像的，有点类似于在${\mathbf{M}}_i$的基础上进行了一个膨胀操作。具体来说，${\mathbf{M}}_i$中第$j$个像素的值${\mathbf{M}}_i^v$计算方法如下：$${\mathbf{M}}_i^v=\{\begin{array}{l}1,D^v\left({\mathbf{Y}}_i\right)\le \lambda \\ 0,\text{ otherwise }\end{array}$$$$D^v\left({\mathbf{Y}}_i\right)=\{\begin{array}{rl}-d\left(v,\partial {\mathbf{Y}}_i\right),& \text{ if }{\mathbf{Y}}_i^v=1\\ d\left(v,\partial {\mathbf{Y}}_i\right),& \text{ if }{\mathbf{Y}}_i^v=0\end{array}$$ 这个$d(v,\partial {\mathbf{Y}}_i)$指的是当前像素$v$与"病变边界"$\partial {\mathbf{Y}}_i$的距离。可以看到，如果$v$本身就在病变区域里面，那么$D^v\left({\mathbf{Y}}_i\right)$就直接给了个负数，从而保证其有更大的概率被选中(即小于$\lambda$)；而如果在病变区域外，我们就认为越近越应该被选中。值得注意的是，这个$\lambda$是可正可负的，从而保证对病变区域进行"膨胀"或者"收缩"。$\lambda$的具体计算过程较为复杂，感兴趣的读者可以阅读原文。