Ускоренная нормализация изображения (массив numpy)

Question

У меня есть функция, которая принимает какое-то изображение на вход и изменяет его значения (масштабы) таким образом, что среднее значение будет 96. Вот функция:

def normalize_image(image: np.ndarray):
    image_median = np.median(image[image > 0])
    image = image * 96.0 / image_median
    image[image > 255] = 255
    return image

Я использую python 3.5.3 и numpy 1.15.2. Я профилировал свой код с помощью cProfile, и оказалось, что эта функция занимает 6% времени (в некоторых случаях до 25% времени), имея всего 50 вызовов. Эти массивы имеют форму (155,256,256).

Я не очень разбираюсь в оптимизации Python, и мне интересно, можно ли это как-нибудь сделать быстрее?

Обычно я бы начал с оптимизации SIMD, но numpy уже активно их использует.

Answer 1

Здесь мало что можно сделать быстро.

Вы уже используете NumPy в векторизованном виде, поэтому выполняется внутренний код C, который, вероятно, уже довольно оптимизирован.

Вычисление медианы может занять гораздо больше времени, чем вычисление среднего значения (поскольку сортировка выполняется). Попробуйте заменить его.

Добавление некоторых скобок должно сохранить разделение массива

image = image * (96.0 / image_median)

потому что между операторами одинакового приоритета Python идет слева направо.