一、问题描述
今天在学习语义分割的代码时, 遇到一个有趣的算法小问题, 本身实现并不复杂, 但不同方法得到的时间复杂度相去甚远, 故通过本文将探索过程详细记录下来。希望对日后遇到的相似问题有所启发。
简要阐述一下问题: 众所周知, 语义分割是一个像素级的分类任务, 需要将输入图像的每一个像素划分到指定的类别。由于是逐像素分类, 故给定的标签常常是和输入图像尺寸一致的Colormap, 其中每一种颜色代表一个特定类别, 比如下面的例子:
如上所述, 每个类别对应一个颜色, 其映射关系是在数据标注阶段设定的。比如上面右图中, 飞机上的每一个像素其像素值均为[128, 0, 0], 行人则用[192, 128, 128]表示。
VOC_COLORMAP = [[0, 0, 0], [128, 0, 0], [0, 128, 0], [128, 128, 0],
[0, 0, 128], [128, 0, 128], [0, 128, 128], [128, 128, 128],
[64, 0, 0], [192, 0, 0], [64, 128, 0], [192, 128, 0],
[64, 0, 128], [192, 0, 128], [64, 128, 128], [192, 128, 128],
[0, 64, 0], [128, 64, 0], [0, 192, 0], [128, 192, 0],
[0, 64, 128]]
VOC_CLASSES = ['background', 'aeroplane', 'bicycle', 'bird', 'boat',
'bottle', 'bus', 'car', 'cat', 'chair', 'cow',
'diningtable', 'dog', 'horse', 'motorbike', 'person',
'potted plant', 'sheep', 'sofa', 'train', 'tv/monitor']
有了这个映射关系, 我们就能将每个像素对应的物体标号找出来。为了方便起见我们可以将类别标签VOC_CLASSES转换成[0, 1, 2, 3, 4,..., 19]表示的数值形式, 现在还有一个问题:
Colormap形式的标签是一个RGB图像, 而根据我们的任务, 一个更理想的标签形式是一张与原图尺寸一致的2-D图像, 其每个像素值等于该像素的类别标签, 即[0, 1, 2,..., 19]中的一个数字。
所以问题就是: 将RGB形式的COLORMAP图像, 转换成2-D形式的LabelMap。
二、解决方法
这部分记录几个解决方法, 及其在单位图像上的耗时。
我在这个小练习中用的数据是自己的数据, 不是VOC数据集; 刚开始实验是用MXNet的NDArray写的, 在函数中加了.asnumpy(), 就发现有的方法非常慢...后来才意识到, 因为nd转np本身就很慢。后来改成numpy的代码, 写博客也方便一些。(顺便吐槽下NDArray对于list的支持太差, 不支持nd.array(list)和ndarray.tolist()等常用方法...)
2.1 暴力解法
首先, 一种不动脑子的做法是遍历输入图像的每一个像素, 返回颜色对应的索引...
首先定义一个映射函数:
def color2index(color):
"""
color: List of [R, G, B]
"""
if color in VerteColorMap:
return VerteColorMap.index(color)
else:
return 0
然后对输入图像的每一个像素进行处理...
def proc_one_img(img):
h, w = img.shape[:2]
idx = np.zeros((h, w))
for i in range(h):
for j in range(w):
idx[i, j]=color2index(img[i, j].tolist())
return idx
运行时间记录:
注意, 上面的color2index函数中第一行的in和第二行的.index可能做了重复的工作, 最近刚开始刷LeetCode, 也是经常这么写。一种可能的优化方法是用try-except代替if, 代码和结果如下:
2.2 map大法
任何时候涉及到将作用于单个元素的函数扩展到这类元素的iterable对象时, 我都会想到map大法...
使用map的思路很简单, 就是上面已经定义了一个处理单个像素位置(h, w)的函数color2index, 现在想要将该函数同时对整个图像所有像素点的位置进行处理。使用map函数的一个难点在于找到扩展的方式:
首先观察color2index, 其输入为List of [R, G, B], 而一张普通图像'img'的shape是[H, W, C], 也就是img[i, j]的结果才是一个[R, G, B]的list。
这里给出自己的经验: 比如刚开始我们不知道要把img弄成怎样的形式才能用map, 那就先假设转换后得到一个iterable对象, 则最后调用的命令为: map(color2index, iterable), 其中iterable是img调整过之后的形式, 那么我们可以用下面的代码检查该iterable是否为正确形式:
for item in iterable:
lbl_value = color2index(item)
显然, 针对这个例子, iterable的形式为: [HxW, C]
def color2index(color):
"""
color: List of [R, G, B]
"""
try:
return VerteColorMap.index(color.tolist())
except ValueError:
return 0
def proc_one_img_v2(img):
"""
img: Np.array with shape HxWxC
"""
h, w = img.shape[:2]
mask = np.array(list(map(color2index, img.reshape(-1, 3)))).reshape(h, w)
return mask
2.3 最优解: 矩阵索引
假如我们能够定义一个矩阵, 用其下标位置代表Uint8空间所有可能的颜色值(256x256x256种颜色), 每一个可能颜色值对应一个标签值(0, 1, 2...), 这样要找出某一个颜色值对应的标签, 只需要用上述矩阵对该颜色索引即可; 同样, 要找出一张图像中所有颜色对应的标签, 直接用该矩阵对该图像(对应的颜色值矩阵)进行索引即可。 掌握矩阵对子下标矩阵的索引是一个远比看起来要困难的任务, 下面举一个简单的例子作为示意。
为了使该映射矩阵形式更简单, 我们可以将传统的[R, G, B]颜色表示方法调整下, 改为用(Rx256x256 + Gx256 + B)一个整数表达的形式。注意这里不需要担心多个不同的[R,G,B]对应同一个整数, 也就是说这种映射关系是一一对应的。
### 现在我们来定义上述映射矩阵: 注意由于256**3数值较大, 需使用32位整数
mapMatrix = np.zeros(256*256*256, dtype=np.int32)
for i, cm in enumerate(VerteColorMap):
mapMatrix[cm[0]*65536 + cm[1]*256 + cm[2]] = i
def SegColor2Label(img, mapMatrix):
"""
img: Shape [h, w, 3]
mapMatrix: color-> label mapping matrix,
覆盖了Uint8 RGB空间所有256x256x256种颜色对应的label
return: labelMatrix: Shape [h, w], 像素值即类别标签
"""
data = img.astype('int32')
indices = data[:,:,0]*65536 + data[:,:,1]*256 + data[:,:,2]
return mapMatrix[indices]
运行时间记录:
我担心这个时间测的不准, 特地把循环次数调成10000次, 还把小数点往后挪了点才找到有效数字...可以看到, 最优方法的速度是方法一的110倍....
当然, 我也验证了下, 三种方法求出的矩阵mask1和mask2, mask3是一样的, 截图就不放了。
三、一点感想
首先就是一个直观感受: 程序员了解的知识哪怕多一点点, 运用到程序中都可以带来明显的优化效果。如果啥都不知道, 那就只能用暴力法遍历...而在图像领域, 逐像素遍历很浮夸的行为, 一定要尽量避免。
另外就是, 矩阵索引确实是非常强大的武器, 这个能力还需要长期的积累修炼才行。map在可迭代对象上也是一如既往的好用, 但是map感觉只是对循环过程提出一点优化, 真正质的提高还是需要更好的算法, 而更好的算法来源于知识的不断积累。
TODO: 分析几种方法的时间复杂度。
