简介
本文的目的是介绍霄龙新开发的用于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时,左侧有两个区域与边界联通,进而后续被面积阈值过滤掉了。
面积阈值
对于集水区的面积,集水区应该面积较小,理想情况下汇集在一点,即周围的水都往一点流,如果本来区域是狭长的,那么就没法汇集于一点,但总归面积不能太大。
因此,稳定集水区面积不得大于给定的阈值。
容积阈值
对于集水区的容积,即集水量,集水量应该比较大,这意味着集水前该集水点所覆盖的面积比较大。
因此,稳定集水区所对应的集水前的面积应该大于容积阈值。