有没有更快的方法将 2d numpy 数组保存到 png

Question

您好，我是 python 新手，我正在尝试将 2d numpy 数组保存到 png 文件中。

我的 2d numpy 数组中的每个元素都是 0 ~ 100 之间的整数，我有一个 getColor() 函数将其映射为 rgb 值。我现在正在做的方式是构建一个与我的 2d numpy 数组具有相同形状的 3 通道 numpy 数组，并将每个值映射到相应的 rgb 值。然而，这需要很多时间，我觉得应该有一个更有效的方法来做到这一点。我的代码目前处理一张图像大约需要 5 秒。

import numpy as np
import imageio

flt_m = get2dArray() # returns a (880*880) numpy array

def getColor(value):
    if(value < 0):
        return (0,0,0)
    elif(value < 50):
        return (100,150,200)
    else:
        return (255,255,255)

canvas = np.zeros((flt_m.shape[0], flt_m.shape[1], 3)).astype(np.uint8)
for row in range(flt_m.shape[0]):
    for col in range(flt_m.shape[1]):
        rgb = getColor(flt_m[row, col])
        for i in range(3):
            canvas[row, col, i] = rgb[i]

imageio.imwrite('test.png', canvas) # saves file to png

Answer 1

@SpghttCd 的回答您已经有了一个很好的解决方案，但是您的写入时间似乎很慢，所以我想了一个替代解决方案...

由于您的图像中只有 2-3 种颜色，您可以编写一个调色板图像（最多支持 256 种颜色），这将占用更少的内存、更少的处理和更少的磁盘空间。它不是为每个像素存储 3 个字节（1 个红色，1 个绿色和 1 个蓝色），而是在每个像素存储一个字节，该字节是 256 色 RGB 查找表或调色板的索引。

import numpy as np
from PIL import Image

# Generate synthetic image of same size with random numbers under 256
flt_im = np.random.randint(0,256,(880,880), dtype=np.uint8)

# Make numpy array into image without allocating any more memory
p = Image.fromarray(flt_im, mode='L')

# Create a palette with 256 colours - first 50 are your blueish colour, rest are white
palette = 50*[100,150,200] +  206*[255,255,255]

# Put palette into image and save
p.putpalette(palette)
p.save('result.png')

显然我无法检查您机器上的性能，但如果我将我的调色板版本与 SpghttCd 的版本进行比较，我会发现速度相差 50 倍：

def SpghttCd(flt_im):
    canvas = np.ones([880, 880, 3], dtype=np.uint8) * 255

    canvas[flt_im<50] = (100, 150, 200)
    imageio.imwrite('SpghttCd.png', canvas)


def me(flt_im):
    # Make numpy array into image without allocating any more memory
    p = Image.fromarray(flt_im, mode='L')

    # Create a palette with 256 colours - first 50 are your blueish colour, rest are white
    palette = 50*[100,150,200] +  206*[255,255,255]

    # Put palette into image and save
    p.putpalette(palette)
    p.save('result.png')

# Generate random data to test with - same for both
flt_im = np.random.randint(0,256,(880,880), dtype=np.uint8)

%timeit me(flt_im)
In [34]: %timeit me(flt_im)                                                                         
34.1 ms ± 1.06 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

In [37]: %timeit SpghttCd(flt_im)                                                                   
1.68 s ± 7.17 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

我注意到，从 PNG 更改为 GIF（对于这种类型的东西同样适用）会导致速度进一步提高 7 倍，即 5 毫秒而不是 34 毫秒。

Answer 2

您可以使用boolean indexing on numpy arrays 设置数组的不同子集。

所以也许你使用：

canvas = np.ones([880, 880, 3], dtype=np.uint8) * 255   # initialize the whole RGB-Array with (255, 255, 255)

canvas[flt_m<50] = (100, 150, 200)      # set all values where flt_m is <50 to (100, 150, 200)

不过，如果你在 flt_m 中确实有负值，你仍然可以添加

canvas[flt_m<0] = (0, 0, 0)