如何使用 numpy 快速迭代和修改像素数组？

Question

首先，我对 Python 及其库比较陌生。

以下代码的目的是将 HDR 图像转换为 RGBM，详见WebGL Insights Chapter 16。

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 中分别迭代每个像素造成的，对于较大的图像可能需要很长时间（对于 3200x1600 大小的图像大约需要 4:30 分钟）。

我的问题是：有没有更有效的方法来实现我的目标？我简要地研究了 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)

我觉得也清楚一点。