python - 将函数应用于 3D numpy 数组

Question

我有一个来自 Image(PIL/Pillow) 对象的 numpy 3D 数组。

 [[178 214 235]
  [180 215 236]
  [180 215 235]
  ..., 
  [146 173 194]
  [145 172 193]
  [146 173 194]]
 ..., 
 [[126 171 203]
  [125 169 203]
  [128 171 205]
  ..., 
  [157 171 182]
  [144 167 182]
  [131 160 180]]]

图像大小约为 500x500 像素。我需要为每个像素应用两个函数。

1. 将 RGB 转换为 LAB(使用 python-colormath 中的函数)此函数采用一维数组，如 [157, 171, 182] 并返回具有 LAB 颜色的一维数组，例如[53.798345635, -10.358443685, 100.358443685]。
2. 使用 scipy.spatial.cKDTree 从自定义调色板中找到最接近的颜色。

自定义调色板是 kd-tree 。

palette = [[0,0,0], [127,127,127], [255,255,255]] #  or [[0.,0.,0.], [50.,0.,0.], [100.,0.,0.]] for LAB color
tree = scipy.spatial.cKDTree(palette)
def find nearest(pixel):
    distance, result = tree.query(pixel)
    new_pixel = palette[result]
    return new_pixel

有没有比使用 Python 迭代更快的解决方案？例如

for row in array:
    for pixel in row:
        apply_fuction1(pixel) # where pixel is one dimensional array like [157 171 182]
        apply_fuction2(pixel)

UPD1 我不知道我做错了什么，但是:

python3 -mtimeit -s'import test' 'test.find_nearest()' # my variant with 2 loops and Image.putdata()
10 loops, best of 3: 3.35 sec per loop
python3 -mtimeit -s'import test' 'test.find_nearest_with_map()' # list comprehension with map and Image.fromarray() by traceur
10 loops, best of 3: 3.67 sec per loop
python3 -mtimeit -s'import test' 'test.along_axis()' # np.apply_along_axis() and Image.fromarray() by AdrienG
10 loops, best of 3: 5.25 sec per loop

def find_nearest(array=test_array):
    new_image = []
    for row in array:
        for pixel in row:
            distance, result = tree.query(pixel)
            new_pixel = palette[result]
            new_image.append(new_pixel)
    im = Image.new('RGB', (300, 200))
    im.putdata(new_image)


def _find_nearest(pixel):
    distance, result = tree.query(pixel)
    new_pixel = palette[result]
    return new_pixel


def along_axis(array=test_array):
    array = np.apply_along_axis(_find_nearest, 2, array)
    im = Image.fromarray(np.uint8(array))


def find_nearest_with_map(array=test_array):
    array = [list(map(_find_nearest, row)) for row in array]
    im = Image.fromarray(np.uint8(array))

Answer 1

抱歉之前的回答，

使用numpy.apply_along_axis

a = np.arange(12).reshape((4,3))
def sum(array):
    return np.sum(array)

np.apply_along_axis(sum, 1, a)
>>> array([ 3, 12, 21, 30])