数字旗手

电气化、自动化、数字化、智能化、智慧化

0%

胞状物体通用分割算法Cellpose解析:使用篇

Cellpose是一个对于胞状物体(比如细胞、晶粒、核、砖块等)进行分割的非常优秀的通用算法,其体现了深度学习在分割这类物体时的强大能力,同时其泛化效果也远超过传统图像处理算法,展现了数据驱动的深度学习所特有的“暴力美学”。

试用

Cellpose的源代码见这里
同时开发者还搭建了网站来方便用户试用Cellpose:
Cellpose快速体验网站:用户可以直接上传自己的图像来直接调用Cellpose,第一时间获得Cellpose的处理效果。
如果用户觉得好,那么可以接着往下深度体验或钻研Cellpose了。

安装

Cellpose的安装有多种方式:

Google Colab在线运行

开发者提供了一个运行在Google Colab上的运行示例脚本,用户可以直接拷贝一份这个脚本到自己的Colab上,然后在线运行。
这种方式的优缺点如下:
优点:可以白嫖Google的算力,不用自己费劲在本地搭建环境及购买硬件;
缺点:Colab有运行时间和资源,且其不支持运行Cellpose的图形交互界面。

可直接执行的二进制程序

开发者使用PyInstaller在Intel处理器上对源代码进行了打包,形成了一个可执行的二进制程序(有Intel MKL加速,但没有GPU支持)。
适用于Windows 10操作系统的程序从这里下载。下载后的程序就是传统的exe程序,可以双击运行它来启动GUI。也有命令行模式:

1
cellpose.exe --dir Pictures/ --chan 2 --save_png

这种方式的优缺点如下:
优点:直接运行开发者打包好的程序,无需自己配置本地环境;
缺点:无法调用GPU计算,计算速度受限;程序启动速度慢;无法自己训练模型,只能使用已有算法模型。

pip包安装

开发者也在pip仓库中上传了Cellpose代码,因此可以使用pip包管理方式来直接安装Cellpose包。
这里还有四种不同的安装方式:分别取决于是否安装GUI、是否支持GPU计算。

第一种:使用CPU计算且无GUI,则:

1
pip install cellpose

第二种:使用CPU计算且有GUI,则:

1
pip install cellpose[gui]

第三种和第四种都是使用GPU计算,因此需要提前配置好GPU环境,即三步走:安装最新GPU驱动、安装CUDA、安装cuDNN,这三步对于深度学习框架都是通用的,可以从搜索引擎上直接搜索教程。
以上依赖安装时注意CUDA版本一定要与mxnet对应好,所以最好是先确定mxnet所需的CUDA版本,然后再具体安装CUDA和cuDNN。
配置好GPU环境后,再安装mxnet的GPU版本:
1
pip install mxnet-cu102

最后按上面第一种或第二种的pip命令来安装cellpose的无GUI版或有GUI版。

源码安装

最自由的方式就是直接从源码安装(虽然从pip包中实质也能获得源码,但pip包有可能不是最新的)。
首先还是配置环境,根据是否要支持GPU,选择是否安装GPU驱动、CUDA和cuDNN。
然后将github源码克隆一下:

1
git clone https://github.com/MouseLand/cellpose.git

进入cellpose文件夹,运行:
1
python -m cellpose

这种方式拥有最大的自由度和灵活性,能训练、能推理、也能自定义代码来满足自己的定制化需求。

上手

上手前准备

对于深度学习来说,它的三大要素是算法、算力和数据。对于Cellpose,关于这三方面:
算力:经过上面的安装过程,算力已经确定,可以是Google Colab的免费算力,也可以是本地环境的CPU或者GPU;
算法:Cellpose的算法模型的基础框架是UNet,具体算法在源码中可以查阅;Cellpose的开发者还提供了其在大量图像上进行训练的预训练模型,该模型会在第一次运行Cellpose时自动从开发者服务器上进行下载;
数据:开发者没有提供其训练数据集,不过提供了16张测试图片,可以从该Google Drive网盘上下载。

GUI模式

终端输入:

1
python -m cellpose

启动GUI。
(1)在GUI中加载图像(拖入图像或从File菜单中加载);
(2)设置模型:Cellpose中有两个模型,cytoplasm和nuclei,即细胞质模型和细胞核模型。比如下面这种图:
cyto-nuclei
根据生物课本上的知识,绿色部分就是细胞质,红色部分是细胞核,根据这两种物质,Cellpose分别有cyto和nuclei两种名称的模型来识别。用户自己的图像有可能不是这种生物图片,但可以根据相似性,来选择用细胞质(细胞质模型其实就是细胞cell模型,即将细胞核也纳入到整个细胞中)还是细胞核模型来分割。
(3)设置通道:选择要分割的图像通道,比如上图中,如果想分割细胞质,即选择green通道;如果想分割细胞核,则选择red通道。如果是分割细胞质,且图中还有细胞核,则将chan设置为细胞质所在通道,而chan2通道设置为细胞核所在通道;如果分割细胞质但里面没有细胞核,则只设置chan即可,chan2设为None;
(4)点击calibrate按钮来预估图中物体的尺寸;也可以手动输入cell diameter来设置。该预估的尺寸会通过左下方的红色圆盘体现;
(5)点击run segmentation来运行模型。当进度条为100%时,模型预测完毕,可以通过是否勾选MASKS ON来调节是否显示分割后的掩膜。
对于上面这张图,如果是分割细胞质/细胞,那么结果为:
cyto
如果是分割细胞核,那么结果为:
nuclei

命令行模式

上面GUI界面中的参数输入同样可以通过命令行模式来实现:
比如:

1
python -m cellpose --dir ~/images_cyto/test/ --pretrained_model cyto --chan 2 --chan2 3 --save_png

还有其他参数可以设置,比如:
dir:图像所在路径;
img_filter:文件名最后的字符(除了扩展名)作为过滤器;
chan:要处理的通道,0是灰度通道,1是red通道,2是green通道,3是blue通道;
chan2:在要处理cyto、同时图中有nuclei时设置,其为nuclei所在通道,0代表None,代表不设置,其他数值所代表的意思同上;
pretrained_model:cyto是细胞质分割模型,nuclei是细胞核分割模型;
diameter:图中物体的平均直径,默认是30;如果设为0,则Cellpose会自动估计;
use_gpu:使用GPU,如果不添加该参数,则使用CPU;
save_png:将分割掩膜存为png,轮廓存为ImageJ所使用的text文件;
save_tif:将分割掩膜存为tif,轮廓存为ImageJ所使用的text文件;
fast_model:通过关闭数据增强以及平均化4 networks来加速代码运行;
all_channels:在所有图像通道上都运行Cellpose,仅用于自定义模型;
no_npy:不存储_seg.npy文件;
batch_size:批处理尺寸。

所有的参数可以通过help参数来查看:

1
python -m cellpose -h

代码模式

与上面两种方式类似,也可以在Python代码中直接调用Cellpose进行编程:

1
2
3
4
5
6
7
8
9
10
11
from cellpose import models
import skimage.io

model = models.Cellpose(gpu=False, model_type='cyto')

files = ['img0.tif', 'img1.tif']

imgs = [skimage.io.imread(f) for f in files]

masks, flows, styles, diams = model.eval(imgs, diameter=None, channels=[0,0],
threshold=0.4, do_3D=False)

可以看出,使用代码调用Cellpose也非常简单,主要就是两步:配置模型models.Cellpose和使用模型进行推理model.eval。

制作数据集

Cellpose的GUI界面不仅能像上面那样用于运行模型,更重要的是可以利用它来制作数据集,从而基于自己的数据来训练模型。
制作数据集的步骤也非常简单:
(1)打开GUI,手动标注物体:右键点击开始标注,再次右键点击或者鼠标回到开始时的圆圈位置则结束标注;(标注时一定将图像中的物体全部都标注完,否则算法会将未标注的物体视为另一类,则会将算法弄晕)
(2)存储标注图像:选择File菜单下的Save masks as PNG,则会将标注图像存为文件;(这个地方有一个坑:此处存储的masks数据格式为np.uint16,如果使用opencv的imread函数读入并显示,会都显示为0;而需要使用skimage的imread函数才能正确读入)
(3)对文件进行组织:将原始图像和标注图像放到一个文件夹内,两者的匹配还需要遵循一定的命名规则,默认为:例如,原始图像名为wells_000.tif,则标注图像需要命名为wells_000_masks.tif。(也可以通过img_filter和mask_filter这两个参数来修改该默认规则)

训练模型

上一步制作好自己的数据集后,可以训练针对该数据集的Cellpose模型。
(在开始训练之前,有一参数需要特别注意,即diameter参数:开发者提供的Cellpose预训练模型中将所有图像进行了resize,使得图像中物体的中位直径都为30像素(细胞质模型)或17像素(细胞核模型);因此如果想训练快速且结果准确,就需要提前将图像resize成其中物体的中位直径约为30像素或17像素;或使用—diameter参数指定图像中大约的中位直径为多少像素。)
训练Cellpose模型可以有两种方式:
(1)在预训练模型基础上进行训练:
这种方式又可以分为两种:
一种是在开发者提供的预训练模型上进行训练:

1
python -m cellpose --train --dir ~/images_cyto/train/ --test_dir ~/images_cyto/test/ --pretrained_model cyto --chan 2 --chan2 1

可以看出,相比于之前的只运行模型,多了—train这个参数及训练集和测试集数据所在路径。
另外一种是在某一给定的预训练模型上进行训练:
1
python -m cellpose --dir ~/images_cyto/test/ --pretrained_model ~/images_cyto/test/model/cellpose_35_0 --save_png

(2)从头训练模型:
1
python -m cellpose --train --dir ~/images_nuclei/train/ --pretrained_model None

即,将参数pretrained_model置为None。

贡献标注数据

上面两步介绍了制作自己的数据集及自己训练模型,实际上开发者还提供了一个额外功能:上传自己的标注数据到开发者服务器上,用于再次训练模型。这样的好处有:(1)对于用户:用户可以不必自己训练模型,等开发者根据用户上传的数据再次训练好模型后,用户就可以直接使用;(2)对于开发者:通过用户“投喂”更多类别的图像,就可以形成更大的数据集来训练Cellpose,从而使得Cellpose的泛化能力和精度都得以提高。
这里有几点注意事项:
(1)先测试一下现有的Cellpose模型在自己的数据上的效果,期间可以尝试更改一下diameter,可能结果会有一点不同;
(2)如果测试效果挺好,即错误较少,那么就没有必要上传这些数据了,因为再次训练的模型性能提高也不大;
(3)如果测试效果很差,那么极有可能自己的数据与Cellpose的训练集中的数据差别很大,那么此时再次基于这些数据的再次训练就有可能有很大提高;
(4)对于上传的数据,物体直径至少有10像素,每张图像上至少有数十个物体:如果图像太小,可以考虑将多张图像拼接起来;如果图像太大,可以考虑将它裁剪成小图,另外,如果图像中有大量非感兴趣的物体,可以将它们直接裁掉;
(5)手动标注时,将物体的边界轮廓勾画出来(对于细胞结构,务必使得轮廓将细胞膜、细胞质和细胞核都包裹进去,这是为了与Cellpose开发者的标注方法一致);
(6)不要直接使用模型预测的结果来上传数据,这会造成“误差”的恶性循环。

论文

创新点

传统的分水岭方法对于有明确边界的对象能够取得较好的分割效果,因为该算法能形成一个个小的“盆地”,这些盆地就代表一个个对象。但是,大多数情形下,不同的对象形成不同“深度”(强度)的盆地,很难统一进行分割。
因此,创建一个关于对象的中间表达intermediate representation,来形成统一的拓扑盆地,就是一个很好的方法。
cellpose做的就是对mask进行模拟扩散simulated diffusion,形成一个矢量场的拓扑映射(以下简称为流场flow field)。

优点

  • 泛化性好,不需要重新训练或调整参数
  • 3D数据可以重用2D的模型,不需要3D的标注(在xyxzyz上都运行2D模型,这样在每个像素上得到6个梯度,然后再平均得到一个完全的3D梯度场。)

网络架构

架构

总体架构使用U-Net
与原U-Net不同的几点有:

  • 跨层连接skip connection时没有使用通道融合,而是直接相加direct summation,目的是为了减少参数量
  • 将其中标准的building blocks换成了残差块residual blocks,其更容易训练,且能使网络加深的同时保证精度
  • 在编码器的末尾使用全局平均池化global average pooling来得到一个图像style的表达,并将该style传入解码器的各个阶段,因为不同的图像风格style应该需要被不同地处理,因此显式地将style作为变量进行考虑。

Ground Truth

Ground Truth包括人工标注的mask及由mask计算的流场的水平和垂直梯度。

由mask计算流场

Cellpose取名来自于OpenPose,该模型会预测出有不同指向的向量场,用来追踪目标。Cellpose中也会预测向量场,其会使得在一个物体内的所有像素都指向该物体的中心。实际上,不要求这些梯度直接指向中心,而是设计成如下:物体内的像素可以按照梯度方向进行移动,经过几轮迭代后,像素都被汇聚到最终无法移动的固定的像素位置,这个位置被称为中心点,所以移动到该位置的像素都被标记为属于该物体。
具体的计算方法是:使用热扩散模拟heat diffusion simulation的方式来生成向量场。

  • 定义物体内最靠近水平和垂直坐标的中值的位置为中心点;
  • 在热扩散模拟中,每次迭代时就在中心点像素上加入一个值为1热源
  • 每次迭代时计算物体内所有像素的温度值,该值等于目标像素的3*3区域(即八邻域)内温度场的平均值;物体外的像素的温度值则统一置为0
  • 上述迭代对于每个物体计算N次,N是物体的水平范围和垂直范围之和的两倍,以此保证热量能扩散到物体的最远的角落;充分迭代后热扩散模拟能达到一个平衡态;
  • 根据最终的温度场分布使用中心差分方法计算梯度向量场,即某像素点上的x方向梯度就是左右两个邻居点的温度的差,y方向梯度就是上下两个邻居的温度的差。

由流场复原mask

神经网络得到的输出是流场的水平梯度、垂直梯度和像素可能性probability,这个输出还不能直接用,需要转换成mask才可以。具体计算方法是:

  • 设定可能性阈值,比如0.5,将像素可能性这个场与该阈值对比,仅考虑可能性大于该阈值的像素点;
  • 在每个像素上,根据两个流场梯度确定的空间差别,做200次的欧拉积分,这样的动力学系统会使得热量反向流动,热量逐渐累积,从而能找到最终的热流汇集处。同时能找到哪些像素点上的热流汇集到了这个最终地点,从而找到属于该热流汇集处所辐射的范围,即找到了物体。

网络训练

Input:单通道灰度图像
Output:流场的水平和垂直梯度 + 像素是物体的可能性
Loss:水平梯度和垂直梯度与GT的两梯度(这里乘以系数5以达到同一量级)之间的MSE + 像素可能性与GTmask之间的交叉熵。

训练配置

  • 特征图大小:第一层的feature map大小为32
  • 迭代数:500
  • 批处理大小:8
  • 优化器:SGD、学习率0.2(最开始的10个迭代中学习率从0线性增加到0.2,以防止初始的不稳定性;从第400个迭代开始,每10个迭代对学习率减半)、动量0.9、权重衰减0.00001

网络推理

对于神经网络output的水平和垂直梯度,综合起来形成流场,以及根据outputInside/outside,对该流场进行进一步的修整;
然后使用梯度追踪gradient tracking算法对该流场重建mask(像素点通过追踪梯度来朝向其最终的不动点,所有收敛到同一不动点的像素被指定为同一个mask)。

数据集

Generalized data

数据具有很大的多样性(细胞、岩石、土豆等),共608张图像,内含有约70000个分割对象。
论文主要的实验结果就是在该数据集上进行训练和评估的。

Nucleus data

作者还准备了一个更大的、专用于Nucleus(细胞核)的数据集,共1139张图像,但数据的多样性明显不如上面的数据集。

数据增强

每一个epoch训练时,数据集及其流场都会被随机增强:

  • 随机缩放:缩放因子在0.51.5之间;
  • 随机旋转:均匀地从0360之间选择旋转角度;
  • 随机平移:平移幅度在(size_original - 224)/2之内,这样保证都是从原始图像内部提取。

增强图像后,以图像中心为中心,以224*224大小进行取样。

效果增强

为了提高模型的预测能力,还使用了一些增强手段:

  • ROI质量估计:计算一个ROI中预测出的流场和从mask中预计算的流场之间的差别来估计ROI的质量;当两者偏差太大时(实测选择的阈值是0.4),舍弃该ROI
  • 模型集成model ensembling:使用同样的模型架构分别训练四次,然后对预测出的流场进行平均;
  • 对测试图像预估物体尺寸并调整大小resizing:对训练集中的图像很容易调整大小到一个统一的尺寸,因为这些图像是提前知道的,但是对于未知的测试图像,事先无法知道图像中物体的平均尺寸。这里使用了一个自动预估图中物体尺寸的方法:(1)从训练集图像的风格style向量中训练一个线性回归模型,这个模型的预测应该在测试集上的效果可以good但不够好perfect;(2)使用上一步线性回归模型给出的尺寸进行调整大小后,基于此,再用Cellpose的分割算法进行细致的尺寸预测。这个预测的尺寸可以作为真实的平均尺寸的参考。
  • 图像分块image tiling:图像分块的原因有两个:(1)使得Cellpose可以运行在任意图像尺寸上;(2)使得Cellpose训练时可以在同样的图像大小(224*224)时进行。具体地,对于图像的两个维度上进行切块时,都保证50%的重叠率overlap(也可以自己设置),此时每个像素都能被计算4次。