python-3.x - 如何遍历数组对每个像素应用阈值

Question

我有一个 3 channel numpy 数组，我想对每个像素应用一个函数。具体来说，我想处理一个图像并返回一个灰度图像，突出显示图像中出现特定颜色的位置。如果红色、绿色、蓝色 channel 与颜色的 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

简单的一行解决方案

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

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

优化的两行解决方案

上述行并不是执行 OP 想要的所有操作的最有效方式。这是一个明显更快的方法(感谢 Paul Panzer 在评论中提出优化建议，不保证可读性):

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 范围内，并在满足条件时返回一个 bool 数组，该数组为 True 和 False否则。

• bool 数组实际上只是一个 0 和 1 的数组，因此我们随后将 bool 数组乘以 255 以获得您想要的最终结果。

优化方案说明

这是使用的优化列表:

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

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

• 我试图找到一种方法来最大程度地减少数组的分配/复制，但这实际上使事情变慢了。这是优化解决方案的一个版本，它恰好分配两个数组(一个用于中间结果，一个用于最终结果)并且不制作其他副本:

  %%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)