简介
本文的目的是介绍霄龙新开发的用于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像素大小的示例图像如下：
（不能直接下载使用，因为这里是为了显示方便，不再是原来的分辨率）
该图的像素矩阵为：
1 | [[ 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0] |
图中有两个白色区域，算法最终目的是找到这两个区域的边界。
生成流场图
如前所述，第一部分的流场仍然是由cellpose计算所得，这是因为cellpose预训练模型的泛化能力很强：
1 | import cv2 |
这里仍然是调用了cellpose的models模块的eval方法。
注意，该eval方法其实是对神经网络计算的一个封装，并不是最原始的纯粹的神经网络计算部分，它默认是将上面的两部分计算都执行的，因此，这里如果单纯的这样调用eval方法，速度仍然会很慢。（最纯粹的神经网络计算是由_run_nets方法执行的，但是它不是一个友好的API，直接调用它的话需要做很多预处理，比如判断通道数、2D或3D，很麻烦。）
好在CellposeModel这个类的eval方法提供了compute_masks这个参数开关，将它置为False的话，就只计算flow，而不计算mask。具体来说，就简单地在这一行改写就行。
这一部分我们只需要flow流场，同时该变量flow是一个复合变量，包含了将流场转为彩色图的值、原始向量场、判断是物体的概率等，这里只需要它的原始向量场，即：
1 | flow[1] |
它的形状是2乘15乘15，2表示向量场有X和Y两个分量，其中第一个分量是Y方向，第二个分量是X方向上。
后面在使用向量场时，会将它的维度调换一下顺序，即：
1 | flow[1].transpose(1,2,0) |
即变为15乘15乘2的矩阵。
后处理加速算法解析
下面就是该加速算法的核心：
1 | water, core, msk = flow2msk(flow[1].transpose(1,2,0), None, 0.1, 5, 10) |
即调用flow2msk函数，其全貌为（注意，以下代码只适用于二维图像，此时可以去源码里看看，应该目前的代码是二维和三维通用了）：
1 | def flow2msk(flow, prob, grad=1.0, area=150, volume=500): |
下面将对该函数逐行理解，其中涉及的算法示意图均来自霄龙分享的笔记，比心~~
（注意，为了方便讲清脉络，下文叙述与源码略有不同，增加了一些冗余代码）
计算流动方向和距离
首先对原像素点进行位置编号：
1 | mark = np.arange(flow.size//2) # 这里除以2是因为flow最后通道有两个向量场 |
示意图如下：
然后计算流场这个向量场的绝对强度：
1 | l = np.linalg.norm(flow, axis=-1) |
其中axis指定为-1，即按最后一个axis来计算2范数，因为flow已经将x和y分量放在了最后一维，所以，上面就是计算了如下公式：
1 | sqrt(x^2+y^2) |
附赠一个理解numpy的axis的博文：
NUMPY AXES EXPLAINED
对原flow场的x和y分量根据上面的绝对强度进行归一化：
1 | flow /= l[:,:,None] |
加上None是为了增加一个维度，与原flow进行统一。
经过上面的归一化操作，现在的flow的x和y分量满足：
1 | x^2+y^2=1 |
接下来根据绝对强度对flow场进行梯度阈值分割，这里的梯度阈值是人为设定的，即grad参数：
1 | flow[l<grad] = 0 |
即判断某像素点上的绝对流场强度，如果强度小于梯度阈值，则直接将该处的流场置为0，即不流动。
（关于该阈值的设置在最后会说明）
进一步地，该算法是模拟水流过程，为了防止水流到区域外面，将第一行和最后一行的Y分量设为0，将第一列和最后一列的X分量设为0：
1 | flow[[0,-1],:,0], flow[:,[0,-1],1] = 0, 0 |
然后挑选出强度特别大的流场所在位置，判断这里的像素点流动的方向：
1 | dn = flow.round().astype(np.int32).reshape(-1,2) |
不过追求速度极致的龙哥嫌numpy的round速度太慢，用以下运算来替代了：
1 | sn = np.sign(flow); sn *= 0.5; flow += sn; |
这个dn的数值有很大讲究，首先其数值为-1、0、1三种，然后它的形状为(225, 2)，225是像素点总个数，2表示两个向量场方向，第一个方向是Y方向，即跨行移动，第二个方向是X方向，即跨列移动，所以如果dn中的某元素为[1, -1]，表示向下移动一行，然后再向左移动一列；0则表示不移动。
根据dn可以知道某处的水是怎样流动的，示意图如下：
具体实操上，还需要更复杂的操作，才能实现示意图中的运算。因为后续会将二维数组压平成一维（这里都压平成一维，是为了适用于任意维度），所以跨行和跨列的移动在内存中的距离是不一样的，比如5乘4的二维矩阵，跨列移动，内存差1，跨行移动则内存差4。
所以首先要根据原始矩阵大小获得其某个像素的邻居像素的标识：
1 | strides = np.cumprod(flow.shape[::-1])//2 |
这里获得邻居标识的过程在find_max和watershed等算法中都用到过，只是这里需要注意因为flow最后一维是2，所以最后除了一个2。
（该方法可以用于任意维度的矩阵的邻居标识的获取）
这里因为是15乘15的矩阵，所以strides的数值为：
1 | [ 1 15 225] |
再将dn的移动方向转换为实际的移动距离：
1 | dn = (dn * strides[-2::-1]).sum(axis=-1) |
对以上代码要分为两步看，首先：
1 | (dn * strides[-2::-1]) |
这是将两个方向上的原来的移动距离转换为真正的移动距离，比如[1, -1]变为了[15, -1]，然后：
1 | (dn * strides[-2::-1]).sum(axis=-1) |
这是对两个方向上的移动进行组合，比如上一步的[15, -1]，就变成了14，所以原来的表示向下移动一行，然后再向左移动一列的[1, -1]数组就变成了一维的移动14距离即可。
有了表征真实移动距离的dn，就可以准确知道水下一刻到达的位置：
该例中，dn的数值为：
1 | dn = [ 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 15 |
第一次的编号转移如下：
1 | rst = np.arange(flow.size//2) + dn |
转移结果为：
1 | rst = [ 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 32 |
位置更新
在上一节的最后，得到了图中各像素的原始编号及下一步要转移到的编号，这就建立了位置a到位置b的映射。
下一步就是要不断迭代，逐步将所有的位置上的编号进行更新。
注意，这里的位置a就是最开始的位置编号，里面的数值是np.range()的顺序编号，比如：
1 | a = np.array([0, 1, 2, 3, 4, 5]) |
而位置b这个变量则有两层意思，举一例子：
1 | b = np.array([1, 2, 3, 4, 5, 5]) |
首先是里面的数值，如上所述，是更新后的位置编号，即a中的0要更新到1上，1要更新到2上，2更新到3，3更新到4，4更新到5，5就指向自己，不更新；
然后是其本身的索引index，即：
1 | b[0], b[1], b[2], b[3], b[4], b[5] |
里面的0到5就暗含了a中的编号，即b中的位置所对应的本来位置。b对本身的索引就代表了一次更新，比如：
1 | b[b] = |
所以，这种不断对自身的索引就是位置更新的机制。
从另外一个角度理解，可以把网格上的编号理解为有向图中的节点（多谢霄龙的指点），比如上面的0号节点指向1号节点，4号指向5号，5号指向自身。这个位置更新的过程就是沿着有向图不断行走的过程。
这里有两种更新方式：一种是单步推演，一种是连锁推演。
单步推演的示意图如下：
即某个位置上的位置更新每次都仅执行一次。
连锁推演的示意图如下：
即某个位置上的位置更新每次可跨越执行。
这两种的原理可通过下面的代码很清楚地理解：
即前一个是以b为被索引对象，而后一个是以不断更新的位置作为被索引对象。可以看出，对于单次推演，本例中需要4次才能稳定；而连锁推演，需要3次即可。
单步推演有线性复杂度，而连锁推演，系统具有对数复杂度。理论可知连锁推演10次，可以将半径扩张到1024，通常这已经足够识别常见的任意物体。
因此，采用连续推演的代码为：
1 | rst = np.arange(flow.size//2) + dn |
最终更新的位置编号为：
1 | rst = [ 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 79 |
获得稳定集水区
经过若干次推演，其实我们是模拟每个网格从最初的位置移动到了新的位置，此时只要对新的位置进行频率统计，就可以得到集水图。
还是一步步看代码怎么实现的。
首先对上面的最新位置编号进行统计，看每个编号出现的次数：
1 | hist = np.bincount(rst, minlength=len(rst)) |
次数统计结果（已经转成了原图像的形状）为：
1 | [[ 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1] |
这个矩阵的数值就代表了集水区，比如0就是水都流走了，是贫水区，1的地方表示水原地没动，数值大于1的地方表示水都汇集在了此处，数值越大表示集水量越大。
下面的工作就是获得这些集水区的区域标记。
首先进行连通域标记：
1 | lab, n = ndimg.label(hist, np.ones((3,3))) |
标记结果为：
1 | lab = |
再统计一下这些区域的面积大小，即每个集水区的面积大小：
1 | areas = np.bincount(lab.ravel()) |
结果为：
1 | areas = [ 89 130 4 2] |
统计这些区域的权重，即每个集水区中集水量的大小，即容积：
1 | weight = ndimg.sum(hist, lab, np.arange(n+1)) |
结果为：
1 | weight = [ 0. 132. 35. 58.] |
因此，集水区的面积和集水量就是衡量这个集水区的两个有效标准。
（1）对于集水区的面积，集水区应该面积较小，理想情况下汇集在一点，即周围的水都往一点流，如果本来区域是狭长的，那么就没法汇集于一点，但总归面积不能太大；
（2）对于集水区的容积，即集水量，集水量应该比较大，这意味着集水前该集水点所覆盖的面积比较大。
算法中提供了面积和容积（集水量）这两个阈值的调控接口，即函数的area和volume参数。
根据面积阈值和体积阈值设定掩膜，这里设置的阈值分别是5和10：
1 | msk = (areas<area) & (weight>volume) |
结果为：
1 | (areas<area) = [False False True True] |
即，对于面积而言，前两个集水区面积太大了，不满足条件；对于容积而言，后三个都满足容积要求，第一个集水区则不满足。
综上，前两个掩膜为False，后两个为True。
根据上面的掩膜创建一个查找表：
1 | lut = np.zeros(n+1, np.int32) |
其值为：
1 | [0 0 1 2] |
有了这个查找表，再将集水区的连通域标记作为索引，就可以得到集水区的区域标记，即：
1 | lut[lab] |
结果为：
1 | [[0 0 0 0 0 0 0 0 0 0 0 0 0 0 0] |
获得集水前区域标记
上面得到了集水区的区域标记，这些标记标明了水最终汇集到的地方，也就是原始网格更新后的位置。
因此，通过集水区的区域标记，再结合之前的位置更新，就可以得到集水前的区域标记（注意是集水“前”）。
代码为：
1 | mask = lut[lab].ravel()[rst] |
最终结果为：
1 | [[0 0 0 0 0 0 0 0 0 0 0 0 0 0 0] |
参数调节
这一节讲一下该加速算法里的三个参数（梯度阈值、面积阈值、容积阈值）的调节原则（所涉图像均为霄龙测试结果）。
梯度阈值
梯度阈值决定集水区的稳定状态，如果阈值设定过小，有可能产生细长区域的过分割，增加阈值会让稳定集水区面积增加，不会打碎较弱的漩涡，而设置过大，有可能会让较弱的流场与背景联通，进而无法被检测到。
梯度阈值0.8时，中间狭长区域汇集成两个核，被分成两份；
梯度阈值1.0时，中间下场区域联通成一个核，是一个整体；
梯度阈值1.3时，左侧有两个区域与边界联通，进而后续被面积阈值过滤掉了。
面积阈值
对于集水区的面积，集水区应该面积较小，理想情况下汇集在一点，即周围的水都往一点流，如果本来区域是狭长的，那么就没法汇集于一点，但总归面积不能太大。
因此，稳定集水区面积不得大于给定的阈值。
容积阈值
对于集水区的容积，即集水量，集水量应该比较大，这意味着集水前该集水点所覆盖的面积比较大。
因此，稳定集水区所对应的集水前的面积应该大于容积阈值。
