有没有更快的方法将二维 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 毫秒。