Lite-UNet: 轻量且高效的细胞定位模型，轻量化的U-Net

论文：Lite-UNet: A Lightweight and Efficent Network for Cell Localization

Paper Link：Lite-UNet: A Lightweight and Efficent Network for Cell Localization

Code： https://github.com/Boli-trainee/Lite-UNet

核心思想：魔改U-Net模型，设计一个轻量高效的细胞定位模型，主要包含3个模块：通过差分卷积增强梯度信息的梯度增强模块（很好用，可以缓解颜色问题）；结合Ghost和CBAM模块的基础模块；利用图神经网络捕获高阶关联的注意力模块。

摘要

细胞定位是病理图像分析领域的一个基本研究领域，其核心目标是精确识别细胞的空间坐标。这项任务一直面临细胞之间的颜色变化大、分布不均匀和边界重叠等挑战。此外，在实际的细胞定位场景中，现有的最先进方法存在高计算成本和缓慢的推理时间，严重降低了计算机辅助的效率。为了解决上述问题，提出了一种轻量级高效的细胞定位模型，名为Lite-UNet。具体来说，Lite-UNet包括三个关键模块。首先，我们引入了基于差异卷积的梯度聚合模块。通过巧妙地利用梯度信息，这个模块有效地减轻了细胞之间广泛的颜色变化带来的挑战。其次，我们提出了一种高效的即插即用的图关联注意模块，通过编码高阶特征关联来优化特征表示能力。最后，我们设计了一个轻量级的Ghost\_CBAM模块，既减轻了细胞分布不均匀的困难，又构成了Lite-UNet的基本模块。广泛的实验证明，我们的Lite-UNet能够快速而准确地定位图像中的细胞，从而进一步提高了计算机辅助医学的效率。

主要贡献点

总之，本文的贡献总结如下：

$\bullet$ 本文提出了一种轻量级和高效的细胞定位模型，以极低的成本实现了竞争性的性能，并可以提高计算机辅助医疗的效率。

$\bullet$ 提出了一种基于差异卷积的梯度聚合模块，有效减轻了细胞之间的大幅度颜色变化问题。

$\bullet$ 设计了基于注意力机制的Ghost\_CBAM模块，有效解决了细胞分布不均匀的问题，同时实现了模型的压缩，基于这一模块构建了轻量高性能的细胞定位模型Lite-UNet。

$\bullet$ 提出了图关联注意模块，编码特征之间的高阶关联，以提供更好的表示。

Introduction

观察细胞对于人类探索微观生物世界是一项至关重要的资源。细胞分析一直是医学图像分析中一个重要领域和具有挑战性的研究课题。在这一领域，细胞定位任务的目的是精确确定图像中每个细胞中心的具体位置。这个任务在许多医学场景中发挥着重要作用，而定位效果将直接影响随后的图像分析。手动在图像中定位细胞是繁琐、耗时和昂贵的工作，因此研究人员一直在致力于开发能够自动完成此任务的算法。然而，细胞图像中的颜色差异大、分布不均匀和边界重叠使得这项工作极具挑战性。

传统的细胞定位方法使用图像处理算法，如阈值处理、Canny边缘检测和颜色识别滤波器来识别细胞。然而，这些方法在细胞重叠或细胞之间颜色差异较大的情况下表现不佳。随着人工智能技术的发展，尤其是深度学习技术的普及，使用深度学习技术直接预测细胞的位置和数量的方法已经取得了令人鼓舞的结果。目前，基于深度学习的细胞定位方法迅速成为诊断和治疗的新辅助手段。许多基于卷积神经网络（CNN）的方法由于CNN强大的非线性拟合能力而实现了出色的细胞定位性能。

目前在细胞定位领域通常采用的范式是回归位置图，然后获取细胞的精确位置。具体而言，首先使用基于CNN的方法将细胞图像映射到相应的位置图，然后对位置图进行后处理以获取关于细胞位置和数量的信息。这里提到的位置图是指可用于细胞定位的地图，如概率图、密度图、方向场图和指数距离变换图等。例如，Xie等人使用全卷积神经网络回归细胞密度图。由于网络采用了全卷积设计，可以在随机尺寸的输入图像上进行预测，从而实现端到端的高效训练。值得注意的是，细胞定位的关键在于生成的位置图的质量，准确反映细胞位置的位置图可以显著改善定位结果。

根据我们的观察，位置图通常难以准确反映细胞的位置信息。其核心原因在于细胞的颜色变化太大，导致位置图中的细胞对应响应的差距太大，最终导致浅色细胞在位置图中的响应极弱，现有的定位模型往往会忽略这些浅色细胞，从而大大降低了模型的定位性能。如图\ref{fig1}所示，当基于普通卷积的模型预测具有大颜色变化的细胞图像时，图像中的浅色细胞在生成的位置图中响应极弱，导致现有的定位模型倾向于忽略这些浅色细胞，这极大地降低了模型的定位性能。值得一提的是，Huang等人\cite{Bcdata}分别标注和预测浅色和深色细胞，巧妙地避免了大颜色变化的困难。然而，这种策略引入了更繁琐的操作和更昂贵的手工标注成本，不利于进一步扩展应用。理想情况下，所有颜色的细胞应该在位置图中对应相同程度的响应。具有大颜色变化的细胞对应于像素值的差异较大，但最终所有细胞需要映射到位置图中的相同响应，从而产生了矛盾。在这方面，我们希望模型能够减小细胞之间的颜色差距，即减小细胞像素的绝对值之间的差距。考虑到差异卷积具有类似的效果，通过在细胞像素上执行局部差异计算，在一定程度上减轻了细胞颜色的差距，从而捕获了相邻像素之间的梯度信息。为此，我们提出了基于差异卷积的梯度聚合模块，通过增强像素之间的梯度信息来缓解这一问题。

目前的细胞定位领域还存在不均匀的细胞分布和高模型计算成本的问题。对于不均匀的细胞分布问题，Guo等人\cite{saunet}使用了自注意力模块来捕捉特征之间的相互关系。然而，自注意力模块通过计算全局特征的关联权重的过程成本较高，而且他们使用了更昂贵的U-Net\cite{unet}作为骨干网络，进一步增加了模型的计算成本。模型的高参数数量和计算工作量限制了其应用，特别是在移动应用和边缘计算部署日益普及的情况下。此外，过慢的推断速度降低了医疗专业人员的体验和工作效率。为了解决这些问题，提出了一种新的Ghost\_CBAM模块，它将计算成本极小的CBAM注意力模块与Ghost模块相结合。该模块在减轻不均匀细胞分布问题的同时，有效实现了模型压缩。本文基于这一模块构建了轻量高性能的细胞定位模型Lite-UNet。

鉴于CNN强大的建模能力，上述基于CNN的模型已经能够实现强大的细胞定位性能。然而，CNN存在两个问题：首先，它通过卷积操作均匀转换图像中的局部特征，忽视了特征之间的相关性。其次，CNN中的感受野受限于局部区域，难以适应场景的复杂拓扑结构。因此，一些研究人员提出了图卷积网络（GCN）来建模场景中的复杂关联。与传统的CNN方法相比，GCN可以编码输入数据的图结构，并从全局图的角度连续学习和汇总相关信息，以更好地表征特征。例如，细胞不均匀分布和各种细胞形状与CNN的感受野不相容，但图卷积网络的拓扑结构可以更好地适应这种非均匀分布。因此，本文提出了一种图相关性注意力模块，通过编码特征之间的高阶关联信息来指导特征的分布，从而提高细胞定位的性能。

总之，本文的贡献总结如下：

$\bullet$ 本文提出了一种轻量级和高效的细胞定位模型，以极低的成本实现了竞争性的性能，并可以提高计算机辅助医疗的效率。

$\bullet$ 提出了一种基于差异卷积的梯度聚合模块，有效减轻了细胞之间的大幅度颜色变化问题。

$\bullet$ 设计了基于注意力机制的Ghost\_CBAM模块，有效解决了细胞分布不均匀的问题，同时实现了模型的压缩，基于这一模块构建了轻量高性能的细胞定位模型Lite-UNet。

$\bullet$ 提出了图关联注意模块，编码特征之间的高阶关联，以提供更好的表示。

Method

本文提出的轻量高效的细胞定位模型 Lite-UNet 如图 \ref{fig2} 所示，展示了每个阶段特征图的维度变化。该模型的输入是细胞图像，输出是相应的细胞位置图。模型由三个主要模块组成：首先，基于差分卷积的梯度聚合 (GA) 模块，作为编码部分的前端，用于充分提取和聚合特征之间的梯度和语义信息；其次，为了解决不均匀的细胞分布问题，同时获得轻量级和高性能模型，我们将网络中的卷积模块替换为 Ghost\_CBAM 模块；最后，图相关性注意力 (GCA) 模块被放置在编码的语义信息丰富端，用于编码特征之间的高阶关联关系。我们将分开解释它们。

Gradient Aggregation (GA) module

如图 \ref{fig1} 所示，细胞分析领域面临的主要挑战是细胞之间的大量颜色差异，这导致模型通常忽视图像中的浅色细胞。具体来说，不同颜色的细胞对应的本地数值在位置图中相似，因此模型难以学习统一的映射表示。从定位结果来看，模型倾向于忽略浅色细胞，导致定位性能下降。因此，我们提出了一种基于差分卷积的梯度聚合模块，通过局部计算细胞像素之间的差异来捕获相邻像素之间的梯度信息，同时减少不同颜色细胞之间像素值的差异。因此，本文提出了一个基于差分卷积的 GA 模块，可以提高模型对梯度信息的利用，缓解细胞之间大量颜色差异的问题。

Attention-based Ghost_CBAM module

图像中的细胞通常分布不均匀，这与大多数方法处理图像中的所有区域一样的策略相矛盾。具体而言，现有的模型不会更加关注细胞密集区域，而且计算成本很高。因此，我们巧妙地将 CBAM（通道注意力模块）和 Ghost（幽灵模块）相结合，得到了 Ghost\_CBAM 模块，它具有轻量和高效的优势。

Graph Correlation Attention (GCA) module

现有的基于关联的图信息嵌入方法的共同范例是以一种聚合方式编码全局特征，然后更新节点特征。然而，在更新过程中存在两个主要问题：1) 具有弱关联的特征会被有选择性地过滤，使得某些本地特征难以更新；2) 优化后的特征直接用于后续操作，可能存在特征过度平滑的问题。因此，本文提出了一个新的 GCA 模块，以更好地捕捉特征之间的关联关系，以更好地指导特征的分布。

Experiments and Analysis

在细胞定位领域，流行的公开可用数据集包括 BCData、Seg_data 和 PSU。BCData 数据集拥有最大的样本量，是本文的主要数据集。

可视化定位效果

Conclusion

本文提出了一种基于U-Net的轻量高效细胞定位模型。我们的Lite-UNet包括三个主要组件：1）梯度聚合模块，可以有效地利用特征的多尺度梯度信息，增强模型对细胞颜色变化的鲁棒性；2）Ghost\_CBAM 模块，可以显著压缩模型的计算成本，而不会丢失大量准确性；3）图相关关注模块，可以通过学习特征之间的高阶关联来优化特征以提高定位性能。全面的实验证明，我们的 Lite-UNet 能够快速准确地定位图像中的细胞，并具有竞争力的性能。与现有模型相比，Lite-UNet 可以部署到计算资源不丰富的更多医疗场景，进一步扩大了计算机医疗辅助的适用范围。在未来，我们将设计一个基于图神经网络的端到端、轻量高性能的细胞定位模型。