Cellpose后处理加速算法解析
source link: https://qixinbo.info/2021/03/03/cellpose-3/
Go to the source link to view the article. You can view the picture content, updated content and better typesetting reading experience. If the link is broken, please click the button below to view the snapshot at that time.
Cellpose后处理加速算法解析
本文的目的是介绍霄龙新开发的用于cellpose后处理加速的算法。
前面已写了两篇文章介绍了cellpose，它是一个非常强大的用于胞状物体分割的算法，具体见：
胞状物体通用分割算法Cellpose解析：使用篇
胞状物体通用分割算法Cellpose解析：开发篇
总体而言，cellpose由两部分组成：第一部分是Unet网络，建立了原图与其流场图flow之间的关系；第二部分是将上一步流场图复原为掩膜mask，即转换为最终的分割结果。
第一部分主要是神经网络的训练和推理，可以跑在CPU或GPU上（当然GPU会更快）；第二部分只能由CPU计算。两者的计算用时见cellpose作者的实测结果：
DNN代表第一部分的神经网络推理部分，实际使用时可以只用1个网络1net，也可用4个网络4net，然后取平均；Postprocessing代表第二部分的后处理部分，即通过流场复原掩膜部分。
可以看出，对于1024乘1024的图像，如果在GPU上只运算1个网络，即1net，第一部分仅用时0.31s，而第二部分却用了6.1s，是第一部分的20倍时长；即使选择CPU，然后运算4个网络4net，第一部分也只用了9.1s，第二部分基本不变，仍为6.1s。
所以，对于cellpose的整体运算，第二部分的由流场复原掩膜的计算成为了速度瓶颈。
然后，毫无意外地，“优化狂魔”霄龙对第二部分下手了，由此有了这个cellpose2msk算法:)。
上链接！！——链接
为了研究一下该算法到底干了啥，构建一个15乘15像素大小的示例图像如下：
（不能直接下载使用，因为这里是为了显示方便，不再是原来的分辨率）
该图的像素矩阵为：
[[ 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]
[ 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]
[ 0 0 255 255 255 0 0 0 0 0 255 255 0 0 0]
[ 0 255 255 255 255 255 0 0 0 0 255 255 255 0 0]
[ 0 255 255 255 255 255 0 0 0 0 0 255 255 0 0]
[ 0 255 255 255 255 255 0 0 0 0 255 255 255 0 0]
[ 0 0 255 255 255 255 0 0 0 0 255 255 255 0 0]
[ 0 0 0 0 0 0 0 0 0 0 255 255 255 0 0]
[ 0 0 0 0 0 0 0 0 0 0 255 255 255 0 0]
[ 0 0 0 0 0 0 0 0 0 0 255 255 255 0 0]
[ 0 0 0 0 0 0 0 0 0 0 255 255 255 255 0]
[ 0 0 0 0 0 0 0 0 0 0 255 255 255 255 0]
[ 0 0 0 0 0 0 0 0 0 0 0 255 255 255 0]
[ 0 0 0 0 0 0 0 0 0 0 0 255 255 0 0]
[ 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]]
图中有两个白色区域，算法最终目的是找到这两个区域的边界。
生成流场图
如前所述，第一部分的流场仍然是由cellpose计算所得，这是因为cellpose预训练模型的泛化能力很强：
import cv2
import cellpose
from cellpose import models, utils
img = cv2.imread("1.png", 0)
use_GPU = models.use_gpu()
model = models.Cellpose(gpu=use_GPU, model_type='cyto')
channels = [0, 0]
_, flow, style, diam = model.eval(
img, diameter=30, rescale=None, channels=channels)
这里仍然是调用了cellpose的models模块的eval方法。
注意，该eval方法其实是对神经网络计算的一个封装，并不是最原始的纯粹的神经网络计算部分，它默认是将上面的两部分计算都执行的，因此，这里如果单纯的这样调用eval方法，速度仍然会很慢。（最纯粹的神经网络计算是由_run_nets方法执行的，但是它不是一个友好的API，直接调用它的话需要做很多预处理，比如判断通道数、2D或3D，很麻烦。）
好在CellposeModel这个类的eval方法提供了compute_masks这个参数开关，将它置为False的话，就只计算flow，而不计算mask。具体来说，就简单地在这一行改写就行。
这一部分我们只需要flow流场，同时该变量flow是一个复合变量，包含了将流场转为彩色图的值、原始向量场、判断是物体的概率等，这里只需要它的原始向量场，即：
flow[1]
它的形状是2乘15乘15，2表示向量场有X和Y两个分量，其中第一个分量是Y方向，第二个分量是X方向上。
后面在使用向量场时，会将它的维度调换一下顺序，即：
flow[1].transpose(1,2,0)
即变为15乘15乘2的矩阵。
后处理加速算法解析
下面就是该加速算法的核心：
water, core, msk = flow2msk(flow[1].transpose(1,2,0), None, 0.1, 5, 10)
即调用flow2msk函数，其全貌为（注意，以下代码只适用于二维图像，此时可以去源码里看看，应该目前的代码是二维和三维通用了）：
def flow2msk(flow, prob, grad=1.0, area=150, volume=500):
l = np.linalg.norm(flow, axis=-1)
flow /= l[:,:,None]; flow[l<grad] = 0
flow[[0,-1],:,0], flow[:,[0,-1],1] = 0, 0
sn = np.sign(flow); sn *= 0.5; flow += sn;
dn = flow.astype(np.int32).reshape(-1,2)
strides = np.cumprod(flow.shape[::-1])//2
dn = (dn * strides[-2::-1]).sum(axis=-1)
rst = np.arange(flow.size//2) + dn
for i in range(10): rst = rst[rst]
hist = np.bincount(rst, minlength=len(rst))
hist.shape = rst.shape = flow.shape[:2]
lab, n = ndimg.label(hist, np.ones((3,3)))
areas = np.bincount(lab.ravel())
weight = ndimg.sum(hist, lab, np.arange(n+1))
msk = (areas<area) & (weight>volume)
lut = np.zeros(n+1, np.int32)
lut[msk] = np.arange(1, msk.sum()+1)
mask = lut[lab].ravel()[rst]
return hist, lut[lab], mask
下面将对该函数逐行理解，其中涉及的算法示意图均来自霄龙分享的笔记，比心~~
（注意，为了方便讲清脉络，下文叙述与源码略有不同，增加了一些冗余代码）
计算流动方向和距离
首先对原像素点进行位置编号：
mark = np.arange(flow.size//2) # 这里除以2是因为flow最后通道有两个向量场
示意图如下：
然后计算流场这个向量场的绝对强度：
l = np.linalg.norm(flow, axis=-1)
其中axis指定为-1，即按最后一个axis来计算2范数，因为flow已经将x和y分量放在了最后一维，所以，上面就是计算了如下公式：
sqrt(x^2+y^2)
附赠一个理解numpy的axis的博文：
NUMPY AXES EXPLAINED
对原flow场的x和y分量根据上面的绝对强度进行归一化：
flow /= l[:,:,None]
加上None是为了增加一个维度，与原flow进行统一。
经过上面的归一化操作，现在的flow的x和y分量满足：
x^2+y^2=1
接下来根据绝对强度对flow场进行梯度阈值分割，这里的梯度阈值是人为设定的，即grad参数：
flow[l<grad] = 0
即判断某像素点上的绝对流场强度，如果强度小于梯度阈值，则直接将该处的流场置为0，即不流动。
（关于该阈值的设置在最后会说明）
进一步地，该算法是模拟水流过程，为了防止水流到区域外面，将第一行和最后一行的Y分量设为0，将第一列和最后一列的X分量设为0：
flow[[0,-1],:,0], flow[:,[0,-1],1] = 0, 0
然后挑选出强度特别大的流场所在位置，判断这里的像素点流动的方向：
dn = flow.round().astype(np.int32).reshape(-1,2)
不过追求速度极致的龙哥嫌numpy的round速度太慢，用以下运算来替代了：
sn = np.sign(flow); sn *= 0.5; flow += sn;
dn = flow.astype(np.int32).reshape(-1,2)
这个dn的数值有很大讲究，首先其数值为-1、0、1三种，然后它的形状为(225, 2)，225是像素点总个数，2表示两个向量场方向，第一个方向是Y方向，即跨行移动，第二个方向是X方向，即跨列移动，所以如果dn中的某元素为[1, -1]，表示向下移动一行，然后再向左移动一列；0则表示不移动。
根据dn可以知道某处的水是怎样流动的，示意图如下：
具体实操上，还需要更复杂的操作，才能实现示意图中的运算。因为后续会将二维数组压平成一维（这里都压平成一维，是为了适用于任意维度），所以跨行和跨列的移动在内存中的距离是不一样的，比如5乘4的二维矩阵，跨列移动，内存差1，跨行移动则内存差4。
所以首先要根据原始矩阵大小获得其某个像素的邻居像素的标识：
strides = np.cumprod(flow.shape[::-1])//2
这里获得邻居标识的过程在find_max和watershed等算法中都用到过，只是这里需要注意因为flow最后一维是2，所以最后除了一个2。
（该方法可以用于任意维度的矩阵的邻居标识的获取）
这里因为是15乘15的矩阵，所以strides的数值为：
[ 1 15 225]
再将dn的移动方向转换为实际的移动距离：
dn = (dn * strides[-2::-1]).sum(axis=-1)
对以上代码要分为两步看，首先：
(dn * strides[-2::-1])
这是将两个方向上的原来的移动距离转换为真正的移动距离，比如[1, -1]变为了[15, -1]，然后：
(dn * strides[-2::-1]).sum(axis=-1)
这是对两个方向上的移动进行组合，比如上一步的[15, -1]，就变成了14，所以原来的表示向下移动一行，然后再向左移动一列的[1, -1]数组就变成了一维的移动14距离即可。
有了表征真实移动距离的dn，就可以准确知道水下一刻到达的位置：
该例中，dn的数值为：
dn = [ 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 15
15 15 0 0 0 0 0 15 15 0 0 0 0 16 15 15 15 15
0 0 0 15 15 15 15 0 0 0 1 16 15 15 14 14 0 0
16 16 15 15 15 0 0 1 1 16 14 -1 -1 0 0 16 16 15
14 15 0 0 1 -14 -14 -16 -1 -1 0 0 1 16 15 14 14 0
0 -14 -14 -15 -15 -16 -16 0 0 1 1 15 14 14 0 0 0 0
-15 -15 -15 0 0 0 1 1 1 -1 -1 0 0 0 0 0 0 0
0 0 0 1 1 -14 -1 -1 0 0 0 0 0 0 0 0 0 0
1 1 -14 -16 -1 0 0 0 0 0 0 0 0 0 0 -14 -14 -15
-16 -1 -15 0 0 0 0 0 0 0 0 0 -14 -14 -15 -15 -16 -15
0 0 0 0 0 0 0 0 0 0 -14 -15 -15 -16 0 0 0 0
0 0 0 0 0 0 0 0 -15 -15 -15 0 0 0 0 0 0 0
0 0 0 0 0 0 0 0 0]
第一次的编号转移如下：
rst = np.arange(flow.size//2) + dn
转移结果为：
rst = [ 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 32
33 34 20 21 22 23 24 40 41 27 28 29 30 47 47 48 49 50
36 37 38 54 55 56 57 43 44 45 47 63 63 64 64 65 52 53
70 71 71 72 73 59 60 62 63 79 78 64 65 67 68 85 86 86
86 88 74 75 77 63 64 63 79 80 82 83 85 101 101 101 102 89
90 77 78 78 79 79 80 97 98 100 101 116 116 117 104 105 106 107
93 94 95 111 112 113 115 116 117 116 117 119 120 121 122 123 124 125
126 127 128 130 131 117 131 132 134 135 136 137 138 139 140 141 142 143
145 146 132 131 147 149 150 151 152 153 154 155 156 157 158 145 146 146
146 162 149 165 166 167 168 169 170 171 172 173 160 161 161 162 162 164
180 181 182 183 184 185 186 187 188 189 176 176 177 177 194 195 196 197
198 199 200 201 202 203 204 205 191 192 193 209 210 211 212 213 214 215
216 217 218 219 220 221 222 223 224]
在上一节的最后，得到了图中各像素的原始编号及下一步要转移到的编号，这就建立了位置a到位置b的映射。
下一步就是要不断迭代，逐步将所有的位置上的编号进行更新。
注意，这里的位置a就是最开始的位置编号，里面的数值是np.range()的顺序编号，比如：
a = np.array([0, 1, 2, 3, 4, 5])
而位置b这个变量则有两层意思，举一例子：
b = np.array([1, 2, 3, 4, 5, 5])
首先是里面的数值，如上所述，是更新后的位置编号，即a中的0要更新到1上，1要更新到2上，2更新到3，3更新到4，4更新到5，5就指向自己，不更新；
然后是其本身的索引index，即：
b[0], b[1], b[2], b[3], b[4], b[5]
里面的0到5就暗含了a中的编号，即b中的位置所对应的本来位置。b对本身的索引就代表了一次更新，比如：
b[b] =
array([2, 3, 4, 5, 5, 5])
所以，这种不断对自身的索引就是位置更新的机制。
从另外一个角度理解，可以把网格上的编号理解为有向图中的节点（多谢霄龙的指点），比如上面的0号节点指向1号节点，4号指向5号，5号指向自身。这个位置更新的过程就是沿着有向图不断行走的过程。
这里有两种更新方式：一种是单步推演，一种是连锁推演。
单步推演的示意图如下：
即某个位置上的位置更新每次都仅执行一次。
连锁推演的示意图如下：
即某个位置上的位置更新每次可跨越执行。
这两种的原理可通过下面的代码很清楚地理解：
即前一个是以b为被索引对象，而后一个是以不断更新的位置作为被索引对象。可以看出，对于单次推演，本例中需要4次才能稳定；而连锁推演，需要3次即可。
单步推演有线性复杂度，而连锁推演，系统具有对数复杂度。理论可知连锁推演10次，可以将半径扩张到1024，通常这已经足够识别常见的任意物体。
因此，采用连续推演的代码为：
rst = np.arange(flow.size//2) + dn
for i in range(10):
rst = rst[rst]
最终更新的位置编号为：
rst = [ 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 79
79 78 20 21 22 23 24 116 116 27 28 29 30 63 63 63 64 64
36 37 38 117 117 117 117 43 44 45 63 79 79 78 78 64 52 53
116 116 116 116 116 59 60 63 79 63 64 78 64 67 68 117 117 117
117 117 74 75 63 79 78 79 63 79 82 83 117 116 116 116 116 89
90 63 64 64 63 63 79 97 98 117 116 117 117 116 104 105 106 107
78 79 79 111 112 113 116 117 116 117 116 119 120 121 122 123 124 125
126 127 128 116 117 116 117 116 134 135 136 137 138 139 140 141 142 143
116 117 116 117 116 149 150 151 152 153 154 155 156 157 158 116 117 117
117 116 149 165 166 167 168 169 170 171 172 173 116 116 116 116 116 149
180 181 182 183 184 185 186 187 188 189 117 117 117 117 194 195 196 197
198 199 200 201 202 203 204 205 116 116 116 209 210 211 212 213 214 215
216 217 218 219 220 221 222 223 224]
获得稳定集水区
经过若干次推演，其实我们是模拟每个网格从最初的位置移动到了新的位置，此时只要对新的位置进行频率统计，就可以得到集水图。
还是一步步看代码怎么实现的。
首先对上面的最新位置编号进行统计，看每个编号出现的次数：
hist = np.bincount(rst, minlength=len(rst))
hist.shape = rst.shape = flow.shape[:2]
次数统计结果（已经转成了原图像的形状）为：
[[ 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1]
[ 1 1 0 0 0 1 1 1 1 1 0 0 1 1 1]
[ 1 0 0 0 0 0 1 1 1 0 0 0 0 1 1]
[ 1 0 0 0 0 0 0 1 1 0 0 0 0 0 1]
[ 1 0 0 11 7 0 0 1 1 0 0 0 0 0 1]
[ 1 0 0 6 11 0 0 1 1 0 0 0 0 0 1]
[ 1 0 0 0 0 0 0 1 1 0 0 0 0 0 1]
[ 1 1 1 0 0 0 1 1 1 0 0 32 26 0 1]
[ 1 1 1 1 1 1 1 1 1 0 0 0 0 0 1]
[ 1 1 1 1 1 1 1 1 1 0 0 0 0 0 3]
[ 1 1 1 1 1 1 1 1 1 0 0 0 0 0 0]
[ 1 1 1 1 1 1 1 1 1 0 0 0 0 0 0]
[ 1 1 1 1 1 1 1 1 1 1 0 0 0 0 1]
[ 1 1 1 1 1 1 1 1 1 1 1 0 0 0 1]
[ 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1]]
这个矩阵的数值就代表了集水区，比如0就是水都流走了，是贫水区，1的地方表示水原地没动，数值大于1的地方表示水都汇集在了此处，数值越大表示集水量越大。
下面的工作就是获得这些集水区的区域标记。
首先进行连通域标记：
lab, n = ndimg.label(hist, np.ones((3,3)))
标记结果为：
lab =
[[1 1 1 1 1 1 1 1 1 1 1 1 1 1 1]
[1 1 0 0 0 1 1 1 1 1 0 0 1 1 1]
[1 0 0 0 0 0 1 1 1 0 0 0 0 1 1]
[1 0 0 0 0 0 0 1 1 0 0 0 0 0 1]
[1 0 0 2 2 0 0 1 1 0 0 0 0 0 1]
[1 0 0 2 2 0 0 1 1 0 0 0 0 0 1]
[1 0 0 0 0 0 0 1 1 0 0 0 0 0 1]
[1 1 1 0 0 0 1 1 1 0 0 3 3 0 1]
[1 1 1 1 1 1 1 1 1 0 0 0 0 0 1]
[1 1 1 1 1 1 1 1 1 0 0 0 0 0 1]
[1 1 1 1 1 1 1 1 1 0 0 0 0 0 0]
[1 1 1 1 1 1 1 1 1 0 0 0 0 0 0]
[1 1 1 1 1 1 1 1 1 1 0 0 0 0 1]
[1 1 1 1 1 1 1 1 1 1 1 0 0 0 1]
[1 1 1 1 1 1 1 1 1 1 1 1 1 1 1]]
再统计一下这些区域的面积大小，即每个集水区的面积大小：
areas = np.bincount(lab.ravel())
areas = [ 89 130 4 2]
统计这些区域的权重，即每个集水区中集水量的大小，即容积：
weight = ndimg.sum(hist, lab, np.arange(n+1))
weight = [ 0. 132. 35. 58.]
因此，集水区的面积和集水量就是衡量这个集水区的两个有效标准。
（1）对于集水区的面积，集水区应该面积较小，理想情况下汇集在一点，即周围的水都往一点流，如果本来区域是狭长的，那么就没法汇集于一点，但总归面积不能太大；
（2）对于集水区的容积，即集水量，集水量应该比较大，这意味着集水前该集水点所覆盖的面积比较大。
算法中提供了面积和容积（集水量）这两个阈值的调控接口，即函数的area和volume参数。
根据面积阈值和体积阈值设定掩膜，这里设置的阈值分别是5和10：
msk = (areas<area) & (weight>volume)
(areas<area) = [False False True True]
(weight>volume) = [False True True True]
msk = [False False True True]
即，对于面积而言，前两个集水区面积太大了，不满足条件；对于容积而言，后三个都满足容积要求，第一个集水区则不满足。
综上，前两个掩膜为False，后两个为True。
根据上面的掩膜创建一个查找表：
lut = np.zeros(n+1, np.int32)
lut[msk] = np.arange(1, msk.sum()+1)
[0 0 1 2]
有了这个查找表，再将集水区的连通域标记作为索引，就可以得到集水区的区域标记，即：
lut[lab]
[[0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]
[0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]
[0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]
[0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]
[0 0 0 1 1 0 0 0 0 0 0 0 0 0 0]
[0 0 0 1 1 0 0 0 0 0 0 0 0 0 0]
[0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]
[0 0 0 0 0 0 0 0 0 0 0 2 2 0 0]
[0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]
[0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]
[0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]
[0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]
[0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]
[0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]
[0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]]
获得集水前区域标记
上面得到了集水区的区域标记，这些标记标明了水最终汇集到的地方，也就是原始网格更新后的位置。
因此，通过集水区的区域标记，再结合之前的位置更新，就可以得到集水前的区域标记（注意是集水“前”）。
代码为：
mask = lut[lab].ravel()[rst]
最终结果为：
[[0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]
[0 0 1 1 1 0 0 0 0 0 2 2 0 0 0]
[0 1 1 1 1 1 0 0 0 2 2 2 2 0 0]
[0 1 1 1 1 1 1 0 0 2 2 2 2 2 0]
[0 1 1 1 1 1 1 0 0 2 2 2 2 2 0]
[0 1 1 1 1 1 1 0 0 2 2 2 2 2 0]
[0 1 1 1 1 1 1 0 0 2 2 2 2 2 0]
[0 0 0 1 1 1 0 0 0 2 2 2 2 2 0]
[0 0 0 0 0 0 0 0 0 2 2 2 2 2 0]
[0 0 0 0 0 0 0 0 0 2 2 2 2 2 0]
[0 0 0 0 0 0 0 0 0 2 2 2 2 2 0]
[0 0 0 0 0 0 0 0 0 2 2 2 2 2 0]
[0 0 0 0 0 0 0 0 0 0 2 2 2 2 0]
[0 0 0 0 0 0 0 0 0 0 0 2 2 2 0]
[0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]]
这一节讲一下该加速算法里的三个参数（梯度阈值、面积阈值、容积阈值）的调节原则（所涉图像均为霄龙测试结果）。
梯度阈值决定集水区的稳定状态，如果阈值设定过小，有可能产生细长区域的过分割，增加阈值会让稳定集水区面积增加，不会打碎较弱的漩涡，而设置过大，有可能会让较弱的流场与背景联通，进而无法被检测到。
梯度阈值0.8时，中间狭长区域汇集成两个核，被分成两份；
梯度阈值1.0时，中间下场区域联通成一个核，是一个整体；
梯度阈值1.3时，左侧有两个区域与边界联通，进而后续被面积阈值过滤掉了。
对于集水区的面积，集水区应该面积较小，理想情况下汇集在一点，即周围的水都往一点流，如果本来区域是狭长的，那么就没法汇集于一点，但总归面积不能太大。
因此，稳定集水区面积不得大于给定的阈值。
对于集水区的容积，即集水量，集水量应该比较大，这意味着集水前该集水点所覆盖的面积比较大。
因此，稳定集水区所对应的集水前的面积应该大于容积阈值。
Recommend
About Joyk
Aggregate valuable and interesting links.
Joyk means Joy of geeK