Numpy: векторизация для нескольких значений

Question

Представьте, что у вас есть изображение RGB и вы хотите обработать каждый пиксель:

import numpy as np
image = np.zeros((1024, 1024, 3))

def rgb_to_something(rgb):
    pass

vfunc = np.vectorize(rgb_to_something)
vfunc(image)

vfunc теперь должно получить каждое значение RGB.Проблема в том, что numpy выравнивает массив и функция получает r0, g0, b0, r1, g1, b1, ..., когда она должна получить rgb0, rgb1, rgb2, ....Можно ли это как-то сделать?

http://docs.scipy.org/doc/numpy/reference/generated/numpy.vectorize.html

Может быть, предварительно преобразовав массив numpy в какой-то специальный тип данных?

Например (конечно, не работает):

image = image.astype(np.float32)
import ctypes
RGB = ctypes.c_float * 3
image.astype(RGB)
ValueError: shape mismatch: objects cannot be broadcast to a single shape

Обновление: главная цель - эффективность здесь.Не векторизованная версия может выглядеть просто так:

import numpy as np
image = np.zeros((1024, 1024, 3))
shape = image.shape[0:2]
image = image.reshape((-1, 3))
def rgb_to_something((r, g, b)):
    return r + g + b
transformed_image = np.array([rgb_to_something(rgb) for rgb in image]).reshape(shape)

Answer 1

Простой способ решить эту проблему - передать весь массив в функцию и использовать в ней векторизованные идиомы.В частности, ваш rgb_to_something также может быть записан

def rgb_to_something(pixels):
    return pixels.sum(axis=1)

, что примерно в 15 раз быстрее, чем ваша версия:

In [16]: %timeit np.array([old_rgb_to_something(rgb) for rgb in image]).reshape(shape)
1 loops, best of 3: 3.03 s per loop

In [19]: %timeit image.sum(axis=1).reshape(shape)
1 loops, best of 3: 192 ms per loop

Проблема с np.vectorize заключается в том, что она обязательно вызываетмного накладных расходов при вызове функции Python применительно к большим массивам.

Answer 2

Вы можете использовать Numexpr для некоторых случаев. Например:

import numpy as np
import numexpr
rgb = np.random.rand(3,1000,1000)
r,g,b = rgb

В этом случае, числоxpr в 5 раз быстрее, чем даже «векторизованное» выражение. Но не все функции могут быть написаны таким образом.

%timeit r*2+g*3/b
10 loops, best of 3: 20.8 ms per loop

%timeit numexpr.evaluate("(r*2+g*3) / b")
100 loops, best of 3: 4.2 ms per loop