Как быстро перебрать и изменить массивы пикселей с NumPy?

Question

Во-первых, я относительно новичок в Python и его библиотеках.

Цель следующего кода - преобразовать изображение HDR в RGBM, как подробно описано в главе 16 WebGL Insights.

import argparse
import numpy
import imageio
import math

# Parse arguments
parser = argparse.ArgumentParser(description = 'Convert a HDR image to a 32bit RGBM image.')
parser.add_argument('file', metavar = 'FILE', type = str, help ='Image file to convert')
args = parser.parse_args()

# Load image
image = imageio.imread(args.file)
height = image.shape[0]
width = image.shape[1]

output = numpy.zeros((height, width, 4))

# Convert image
for i in numpy.ndindex(image.shape[:2]):
    rgb = image[i]
    rgba = numpy.zeros(4)
    rgba[0:3] = (1.0 / 7.0) * numpy.sqrt(rgb)
    rgba[3] = max(max(rgba[0], rgba[1]), rgba[2])
    rgba[3] = numpy.clip(rgba[3], 1.0 / 255.0, 1.0)
    rgba[3] = math.ceil(rgba[3] * 255.0) / 255.0
    output[i] = rgba

# Save image to png
imageio.imsave(args.file.split('.')[0] + '_rgbm.png', output)

Код работает и дает правильные результаты, но делает это очень медленно. Это, конечно, вызвано итерацией каждого пикселя в Python, что для больших изображений может занять много времени (около 4:30 минут для изображения размером 3200x1600).

Мой вопрос: есть ли более эффективный способ добиться того, что я хочу? Я кратко рассмотрел векторизацию и трансляцию в numpy, но пока не нашел способа применить их к моей проблеме.

Редактировать:

Благодаря Mateen Ulhaq я нашел решение:

# Convert image
rgb = (1.0 / 7.0) * numpy.sqrt(image)
alpha = numpy.amax(rgb, axis=2)
alpha = numpy.clip(alpha, 1.0 / 255.0, 1.0)
alpha = numpy.ceil(alpha * 255.0) / 255.0
alpha = numpy.reshape(alpha, (height, width, 1))
output = numpy.concatenate((rgb, alpha), axis=2)

Это завершается всего за несколько секунд.

Answer 1

Линия

for i in numpy.ndindex(image.shape[:2]):

просто перебирает каждый пиксель. Вероятно, быстрее избавиться от цикла и обработать каждый пиксель в каждой строке кода ("векторизовано").

rgb = (1.0 / 7.0) * np.sqrt(image)
alpha = np.amax(rgb, axis=2)
alpha = np.clip(alpha, 1.0 / 255.0, 1.0)
alpha = np.ceil(alpha * 255.0) / 255.0
alpha = numpy.reshape(alpha, (height, width, 1))
output = np.concatenate((rgb, alpha), axis=2)

Я думаю, что это также немного яснее.