python - 如何加速查找操作的嵌套循环？

Question

我正在编程 halftoning of images for laser-engraving .在给定的设置下，激光只会打开或关闭，所以我可以给它提供 1 位深度的二进制图像。因此，我将具有 8 位深度(0 到 255)的灰度图像转换为具有 1 位深度(0 到 1)的二值图像。

我以下面的两张图片为例。左边是灰度图像。右边是用 3x3 正方形二进制像素替换每个像素的结果。结果看起来很相似，因为灰色来自黑色像素的密度。

我目前的尝试是使用嵌套循环来访问像素，并将输出图像中的像素替换为字典中的查找值:

import math
import time

import numpy as np

TONES = [[0, 0,
          0, 0],
         [0, 1,
          0, 0],
         [1, 1,
          0, 0],
         [1, 1,
          0, 1],
         [1, 1,
          1, 1]]

def process_tones():
    """Converts the tones above to the right shape."""
    tones_dict = dict()

    for t in TONES:
        brightness = sum(t)
        bitmap_tone = np.reshape(t, (2, 2)) * 255
        tones_dict[brightness] = bitmap_tone
    return(tones_dict)

def halftone(gray, tones_dict):
    """Generate a new image where each pixel is replaced by one with the values in tones_dict.
    """

    num_rows = gray.shape[0]
    num_cols = gray.shape[1]
    num_tones = len(tones_dict)
    tone_width = int(math.sqrt(num_tones - 1))

    output = np.zeros((num_rows * tone_width, num_cols * tone_width),
                         dtype = np.uint8)

    # Go through each pixel
    for i in range(num_rows):
        i_output = range(i * tone_width, (i + 1)* tone_width)

        for j in range(num_cols):
            j_output = range(j * tone_width, (j + 1)* tone_width)

            pixel = gray[i, j]
            brightness = int(round((num_tones - 1) * pixel / 255))

            output[np.ix_(i_output, j_output)] = tones_dict[brightness]

    return output

def generate_gray_image(width = 100, height = 100):
    """Generates a random grayscale image.
    """

    return (np.random.rand(width, height) * 256).astype(np.uint8)

gray = generate_gray_image()
tones_dict = process_tones()

start = time.time()
for i in range(10):
    binary = halftone(gray, tones_dict = tones_dict)
duration = time.time() - start
print("Average loop time: " + str(duration))

结果是:

Average loop time: 3.228989839553833

对于 100x100 的图像，平均循环需要 3 秒，与 OpenCV 的函数相比，这似乎很长

我检查了How to speed-up python nested loop?和 Looping over pixels in an image而且我没有立即看到如何向量化此操作。

我怎样才能加快这个查找操作的嵌套循环？

Answer 1

诀窍是不要像您那样以如此低的粒度进行迭代，而是将大部分工作卸载到优化的 numpy 函数。

---

从概念上讲，我们可以将输出图像视为一组较小的图像(称它们为“ channel ”)，每个图像都保存半色调网格中一个位置的数据。

然后可以通过简单的查找生成单独的 channel 图像，在 Numpy 中我们可以简单地通过 indexing 来完成。带有灰度图像的查找表(即 LUT[image])。

查找表

假设我们按以下方式定义“图 block 大小”(一个半色调图案的大小)和各个色调图 block :

TILE_SIZE = (2, 2) # Rows, Cols

TONES = np.array(
    [[0, 0,
      0, 0],
     [0, 1,
      0, 0],
     [1, 1,
      0, 0],
     [1, 1,
      0, 1],
     [1, 1,
      1, 1]]
    , dtype=np.uint8) * 255

我们首先使用np.linspace计算灰度和色调指数之间的映射。然后对于每个位置，我们根据音调的定义创建查找表(使用查找技术来这样做)。

def generate_LUTs(tones, tile_size):
    num_tones, num_tiles = tones.shape
    tile_rows, tile_cols = tile_size
    assert(num_tiles == (tile_rows * tile_cols))

    # Generate map between grayscale value and tone index
    gray_level = np.linspace(0, (num_tones - 1), 256, dtype=np.float32)
    tone_map = np.uint8(np.round(gray_level))

    # Generate lookup tables for each tile
    LUTs = []
    for tile in range(num_tiles):
        LUTs.append(tones[:,tile][tone_map])

    return LUTs

合并 channel

现在，将 channel 合并成一个完整的输出图像。

第一步是reshape每个 channel 图像，使其只有一列。

然后，我们可以使用 np.hstack 组合共享相同半色调图案行的所有 channel 图像。 .

接下来，我们 reshape 结果，使其具有与输入图像相同的行数(即它们现在的列数是输入图像的两倍)。

我们使用 np.hstack 再次组合所有重新整形的图像.

最后，我们 reshape 结果，使其具有正确的行数(根据图 block 大小)，然后我们就完成了。

在代码中(适用于任何图 block 大小):

def halftone(image, LUTs, tile_size):
    tiles = []
    for tile in range(len(LUTs)):
        tiles.append(LUTs[tile][image])

    image_rows, _ = image.shape
    tile_rows, tile_cols = tile_size

    merged_rows = []
    for row in range(tile_rows):
        row_tiles = tiles[row * tile_cols:(row + 1) * tile_cols]
        merged_row = np.hstack([row_tile.reshape(-1, 1) for row_tile in row_tiles])
        merged_rows.append(merged_row.reshape(image_rows, -1))

    return np.hstack(merged_rows).reshape(image_rows * tile_rows, -1)

---

使用示例:

LUTs = generate_LUTs(TONES, TILE_SIZE)
binary = halftone(gray, LUTs, TILE_SIZE)

---

示例输出:

还有 3x3 block :