如何通过数组对每个像素应用阈值

Question

我有一个 3 通道 numpy 数组，我想对每个像素应用一个函数。具体来说，我想处理图像并返回灰度图像，突出显示图像中出现特定颜色的位置。如果红色、绿色、蓝色通道与颜色的距离在 L2 的 10 以内：(30,70,130)，则将该像素在灰度图像上的值设置为 255，否则为 0。

我目前的做法是：

def L2_dist(p1,p2):
    dist = ( (p1[0]-p2[0] )**2 + (p1[1]-p2[1] )**2 + (p1[2]-p2[2] )**2 ) **0.5
    if dist<10: return 255
    return 0

def colour_img(image):
    colour = my_colour
    img_dim = image.shape
    new_img = np.zeros((img_dim[0],img_dim[1])) # no alpha channel replica
    for c in range(img_dim[0]):
        for r in range(img_dim[1]):
            pixel = image[r,c,:3]
            new_img[r,c] = L2_dist(colour,pixel)
    return new_img

但是速度很慢。我怎样才能更快地做到这一点而不是使用循环？

Answer 1

你可以这样做

color = np.array([30, 70, 130])
L2 = np.sqrt(np.sum((image - color) ** 2, axis=2))  # L2 distance of each pixel from color

img_dim = image.shape
new_img = np.zeros((img_dim[0], img_dim[1]))
new_img[L2 < 10] = 255

但是正如您所看到的，我们对数组进行了两次迭代，首先计算 L2，然后在 L2 < 10 中进行阈值处理，我们可以通过嵌套循环改进它，就像在您的代码中所做的那样。但是，python 中的循环很慢。因此，JIT 编译该函数以获得最快的版本。下面我用numba：

import numba as nb

@nb.njit(cache=True)
def L2_dist(p1,p2):
    dist = (p1[0]-p2[0] )**2 + (p1[1]-p2[1] )**2 + (p1[2]-p2[2] )**2
    if dist < 100: return 255
    return 0

@nb.njit(cache=True)
def color_img(image):
    n_rows, n_cols, _ = image.shape
    new_img = np.zeros((n_rows, n_cols), dtype=np.int32)
    for c in range(n_rows):
        for r in range(n_cols):
            pixel = image[r, c, :3]
            new_img[r,c] = L2_dist(color,pixel)
    return new_img

时间安排：

# @tel's fully optimised solution(using einsum to short circuit np to get to BLAS directly, the no sqrt trick)
128 µs ± 6.94 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

# JITed version without the sqrt trick
30.8 µs ± 10.2 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

# JITed version with the sqrt trick
24.8 µs ± 11.9 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

HTH。

Answer 2

简单的一条线解决方案

你可以像这样在一行中做你想做的事：

new_img = (((image - color)**2).sum(axis=2)**.5 <= 10) * 255

优化的两线方案

上述行并不是执行 OP 想要的所有操作的最有效方式。这是一种明显更快的方法（感谢 Paul Panzer 提出的 cmets 优化建议，不保证可读性）：

d = image - color
new_img = (np.einsum('...i, ...i', d, d) <= 100) * 255

时间安排：

给定一些 100x100 像素的测试数据：

import numpy as np

color = np.array([30, 70, 130])
# random data within [20,60,120]-[40,80,140] for demo purposes
image = np.random.randint(10*2 + 1, size=[100,100,3]) + color - 10

这是 OP 方法的时间安排和此答案中的解决方案的比较。单线解决方案比 OP 快约 100 倍，而完全优化的版本快约 300 倍：

%%timeit
# OP's code
img_dim = image.shape
new_img = np.zeros((img_dim[0],img_dim[1])) # no alpha channel replica
for c in range(img_dim[0]):
    for r in range(img_dim[1]):
        pixel = image[r,c,:3]
        new_img[r,c] = L2_dist(color,pixel)

43.8 ms ± 502 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%%timeit
# one line solution
new_img = (((image - color)**2).sum(axis=2)**.5 <= 10) * 255

439 µs ± 13.7 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%%timeit
# fully optimized solution
d = image - color
new_img = (np.einsum('...i, ...i', d, d) <= 100) * 255

145 µs ± 2.29 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

简单的单行解法说明

作为第一个解决方案给出的简单单行将：

• 求image（形状为(m, n, 3)的数组）和color（形状为(3)的数组）中每个像素之间的欧几里得距离。

• 检查这些距离中的任何一个是否在10 范围内，并在满足条件的地方返回一个布尔数组True，否则返回False。

• 布尔数组实际上只是0s 和1s 的数组，因此我们将布尔数组乘以255 得到您想要的最终结果。

  李>

优化方案说明

以下是使用的优化列表：

• 使用einsum 计算距离计算所需的平方和。在底层，einsum 使用 Numpy 包装的 BLAS 库来计算所需的 sum-product，因此它应该更快。

• 通过比较距离的平方与阈值的平方来跳过平方根。

• 我试图找到一种方法来最小化数组的分配/复制，但这实际上让事情变慢了。这是优化解决方案的一个版本，它恰好分配了两个数组（一个用于中间结果，一个用于最终结果）并且不制作其他副本：

  %%timeit
  # fully optimized solution, makes as few array copies as possible
  scr = image.astype(np.double)
  new_img = np.zeros(image.shape[:2], dtype=np.uint8)
  np.multiply(np.less_equal(np.einsum('...i,...i', np.subtract(image, color, out=scr), scr, out=scr[:,:,0]), 100, out=new_img), 255, out=new_img)
  
  232 µs ± 7.72 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)