图像配准图像融合再探索/图像像素点遍历加速

Posted 2023-03-03 zstar-_

前言

在我先前的博文【图像配准】多图配准/不同特征提取算法/匹配器比较测试中，提到了图像融合的一种方式，相关代码如下：

result[0:imageB.shape[0], 0:imageB.shape[1]] = np.maximum(imageB, result[0:imageB.shape[0], 0:imageB.shape[1]])

该融合思路是两幅图重复的部分取两者的较大值。
不过在后续实验中，遇到的一个问题是，如果某张图上有一些颜色标记，那么在用取大融合之后，标记颜色会发生变化。
因此，更合理的融合方式是重叠部分直接采用一幅图的原图。

实现思路

方式一：纯像素遍历

要实现该功能，最容易想到的就是来个双循环遍历所有像素：

def merge(imageB, result):
    for r in range(imageB.shape[0]):
        for c in range(imageB.shape[1]):
            # print(result[r, c])
            # print(type(result[r, c]))  #  <class 'numpy.ndarray'>
            if result[r, c, 0] == 0:
                result[r, c] = imageB[r, c]
    return result

imageB是需要融合上的图片，result是已根据imageB做仿射变换的图像。这里的思路是根据imageB的范围在result中进行搜寻，如果是黑色部分，则用imageB的位置替代。
注：这里判断黑色稍微取巧了一些，result的第三个维度通常包含(r,g,b)（如果用opencv直接读取则是(b,g,r)），纯黑的r,g,b是(0,0,0)，通常来说，只要一个通道为0就可以认定其为黑色，因为自然图像中很少包含纯色，通常都是三通道混杂的颜色。
如果需要严格判断纯色，可以将判断语句修改为：

if (result[r, c] == np.array((0, 0, 0))).all():
    result[r, c] = imageB[r, c]

采用两张7k x 5k分辨率的图片做配准融合，所耗费的时间为：138.10秒。

方式二：纯像素遍历+GPU

显然，配准两张图片花费2分多种实在是太慢了，遍历像素点的计算太多，CPU效率不够快。那么，是否可以将该部分的计算放到GPU中去进行呢？当然可以，Numba提供非常方便的装饰器。

Numba是Nvidia开发的一个适用于 Python 代码的开源式即时编译器。
Numba的具体介绍可参考官方的介绍文档：https://www.nvidia.cn/glossary/data-science/numba/

使用起来很容易，在原函数上加一个装饰器即可。

from numba import jit

@jit
def merge(imageB, result):
    for r in range(imageB.shape[0]):
        for c in range(imageB.shape[1]):
            # print(result[r, c])
            # print(type(result[r, c]))  #  <class 'numpy.ndarray'>
            if (result[r, c] == np.array((0, 0, 0))).all():
                result[r, c] = imageB[r, c]
    return result

采用两张7k x 5k分辨率的图片做配准融合，所耗费的时间为：20.85秒，提升非常明显。

方法三：np.where

既然遍历像素效率太低，那么可以使用numpy的where函数来搜寻范围，numpy封装了一些并行计算的方式，可以极大提升效率。

rows, cols = np.where(result[:imageB.shape[0], :imageB.shape[1], 0] == 0)
result[rows, cols] = imageB[rows, cols]

使用时仅需两行代码，替代了上述两个复杂的for循环。
采用两张7k x 5k分辨率的图片做配准融合，所耗费的时间为：14.54秒，速度进一步提升。

总结

涉及到搜寻像素范围时，优先使用np.where；遇到密集计算时，可以尝试用@jit进行GPU加速。

完整代码

import pickle
import time
from numba import jit
import cv2
import numpy as np


# 去除图像黑边
def cutBlack(pic):
    rows, cols = np.where(pic[:, :, 0] != 0)
    min_row, max_row = min(rows), max(rows) + 1
    min_col, max_col = min(cols), max(cols) + 1
    return pic[min_row:max_row, min_col:max_col, :]


# merge
@jit
def merge(imageB, result):
    for r in range(imageB.shape[0]):
        for c in range(imageB.shape[1]):
            # print(result[r, c])
            # print(type(result[r, c]))  #  <class 'numpy.ndarray'>
            if (result[r, c] == np.array((0, 0, 0))).all():
                result[r, c] = imageB[r, c]
    return result


if __name__ == '__main__':
    s_t = time.time()
    with open("H.pkl", "rb") as f:
        H = pickle.load(f)
    imageA = cv2.imread("one_line/DSC01026.JPG")
    imageB = cv2.imread("one_line/DSC01027.JPG")
    result = cv2.warpPerspective(imageA, H,
                                 ((imageA.shape[1] + imageB.shape[1]) * 2, (imageA.shape[0] + imageB.shape[0]) * 2))
    # 问题：merge这里叠加-某些颜色变淡
    # result[0:imageB.shape[0], 0:imageB.shape[1]] = np.maximum(imageB, result[0:imageB.shape[0], 0:imageB.shape[1]])
    # 方法一/二：像素点遍历：
    # result = merge(imageB, result)
    # 方法三：np.where
    rows, cols = np.where(result[:imageB.shape[0], :imageB.shape[1], 0] == 0)
    result[rows, cols] = imageB[rows, cols]
    result = cutBlack(result)  # 结果去除黑边
    cv2.imwrite("through.tif", result[:, :, [0, 1, 2]])
    e_t = time.time()
    print("耗时：", e_t - s_t)