NIPS 2022细胞分割赛道第一名方案

NIPS 2022细胞分割赛道第一名方案

简介

NIPS（NeurIPS），全称神经信息处理系统大会(Conference and Workshop on Neural Information Processing Systems)，是一个关于机器学习和计算神经科学的国际会议。

NIPS是机器学习领域的顶级会议，在中国计算机学会的国际学术会议排名中，NIPS为人工智能领域的A类会议。NIPS Competition Track是其官方举办的竞赛。在本届竞赛中，第四范式AutoX团队在细胞图像分割赛道获得了一等奖。

赛题介绍

基于显微镜图像的单细胞分析是目前生命科学和生物学等领域的前沿和热点问题。细胞分割是单细胞分析中关键的环节一环，它的目标是从显微镜图像中分割出所有的细胞。深度学习已广泛用于图像分割，然后实际场景中，手动注释细胞非常耗时且成本高昂，我们很难收集大量的标记数据来训练模型。另一方面，目前已有的数据集通常限于一种模态，缺乏多样性，导致训练模型的泛化性能较差。

在此背景下，主办方设立了本次多模态显微镜图像细胞分割竞赛，主办方联合了全球十多个实验室累计收集了40多个不同生物实验的图像，是目前多样性最高的数据集。同时还提供了部分未标记的数据，鼓励参赛者利用无标记的图像和公开的数据集来提升模型的泛化性能。

1. 数据情况

训练集包含1000张标注好的图像块，以及1500多张无标注的图像。

2. 评测指标

评测指标同时包含分割精度和分割效率。分割精度指标: F1 Score (先统计分割正确细胞的recall和precision，然后计算F1 score，每个细胞跟对应金标准的交并比IoU超过0.5即认为分割正确)。分割效率指标：分割单张图像的时间（显存不超过1500MB）。

3. 竞赛官网

https://neurips22-cellseg.grand-challenge.org/

解决方案

我们提出了一种通用半监督细胞实例分割模型VSM（Versatile Semi-supervised Model），该模型的主要框架图如图1所示。考虑到区别于自然图像，细胞图像往往具有简单的语义信息和固定的结构，使得低级语义信息——如边缘、形状、纹理等在细胞分割问题上十分重要，我们将任务分为两个阶段，即先检测细胞，再在通过边界框裁剪出的图像中分割目标。VSM主要由三部分组成：预处理模块、检测部分、分割部分，另外，为了更好的利用比赛中提供的无标记数据，我们设计了一种基于伪标签学习的半监督训练流程。我们将在之后详述这部分。

VSM流程图

1. 预处理

由于输入的图像来自不同的模态，并处于不同的图像空间以及有着不同的图像格式，我们需要首先需要通过预处理将图像统一到相同的图像空间，预处理包括三部分：通道对齐、强度归一化和滑动裁剪。首先，我们进行通道对齐，以确保所有图像具有相同数量的通道，其中灰度图像的通道将被复制。然后，我们将每个通道的像素值归一化为0到255之间。将省略前1％的像素和最后99％的像素（分别设为0和255）。最后，我们将图像裁剪成具有1536的窗口和1024的步幅的块，以处理大分辨率图像和WSI。在推理和训练中将进行相同的预处理。

2. 检测部分

在对比赛数据集进行初步测试后，我们发现YOLOX可以在不调整超参数的情况下提供出色的性能，此外，YOLOX是一种anchor-free模型，具有多层次检测头，非常适合尺寸差异极大的情况。因此，我们选择YOLOX作为细胞检测部分。为了进一步提高检测质量，我们使用TTA（Test-Time Augmentation）来改进模型，输入图像将通过0、90、180和270度旋转以及镜像旋转来复制，将来自八个方向的图像副本同时进行检测并集成。在开发过程中，我们使用WBF（Weighted Box Fusion）来合并来自不同副本的边界框，并在最终的比赛提交中用NMS替换它，以加速推理阶段。

3. 分割部分

U-Net可以在分割过程中整合低级特征和高级特征，在图像具有简单明了的语义信息的细胞分割任务中提供出色的性能。此外，在我们的实验中，我们发现模型推理速度的瓶颈在于分割部分。作为轻量级模型，U-Net在提高推理效率方面非常有帮助。因此，我们使用U-Net作为分割网络的骨干。具体来说，我们选择所有置信度大于0.5的边界框裁剪框内的图像，并将其调整为64x64大小来作为模型输入，其中非整数坐标的边界框将被舍入到最近的整数。完成分割后，我们将分割结果恢复到原来的尺寸，分配唯一索引并将其粘贴到总分割掩膜中。

4. 基于伪标签学习的半监督训练策略

基于伪标签学习的半监督训练流程如下图所示。

基于伪标签学习的半监督训练流程

首先，我们使用COCO预训练初始化的模型在LIVECell数据集和Sartorius数据集上进行预训练。这里我们还尝试了添加TissueNet数据集对模型进行预训练，但是没有观察到额外的提升。然后，对预训练模型在有标签数据上进行微调，并将其视为初步模型。我们使用它为无标记数据生成伪标签，在此过程中使用TTA来提高伪标签质量。

然后，我们使用包括翻转和旋转在内的多种数据增强来复制有标签数据，并将其与无标签数据及其伪标签混合，其中伪标签被视为硬标签。我们将有标签数据与无标签数据的比例控制在5：1。

最后，我们再次在这个合并后的数据集上微调预训练模型，以获得更好的模型。