Existe-t-il un moyen plus rapide d'enregistrer des tableaux numpy 2d en png ?

Question

Bonjour, je suis nouveau dans le monde de Python et j'essaie d'enregistrer un tableau numpy 2d dans un fichier png.

Chaque élément de mon tableau numpy 2d est un entier entre 0 ~ 100, et j'ai un getColor() pour le convertir en valeurs rgb. Pour l'instant, je construis un tableau numpy à 3 canaux de la même forme que mon tableau numpy 2d, et je convertis chaque valeur en valeur rgb correspondante. Cependant, cela prend beaucoup de temps et je pense qu'il devrait y avoir un moyen beaucoup plus efficace de le faire. Mon code prend actuellement environ 5 secondes pour traiter une image.

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

Vous avez déjà une bonne solution avec la réponse de @SpghttCd, mais vos temps d'écriture semblent très lents, donc j'ai réfléchi à une solution alternative...

Comme vous n'avez que 2 ou 3 couleurs dans votre image, vous pourriez peut-être écrire une image palettisée (qui prend en charge jusqu'à 256 couleurs) et cela devrait prendre moins de mémoire, moins de traitement et moins d'espace disque. Au lieu de stocker 3 octets (1 pour le rouge, 1 pour le vert et 1 pour le bleu) pour chaque pixel, il stocke un seul octet à chaque pixel et cet octet est un index dans une table de consultation RVB de 256 couleurs ou palette.

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')

Je ne peux évidemment pas vérifier les performances sur votre machine, mais si je compare ma version palettisée avec celle de SpghttCd, j'obtiens une différence de vitesse massive de 50x :

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)

Je constate que le passage d'un PNG à un GIF (tout aussi performant pour ce type de chose) entraîne une amélioration supplémentaire de 7x la vitesse, soit 5ms au lieu de 34ms.

Answer 2

Vous pouvez utiliser indexation booléenne sur les tableaux numpy pour définir différents sous-ensembles d'un tableau.

Alors peut-être que vous utilisez :

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)

_Cependant, si vous avez des valeurs négatives dans flt_m vous pouvez toujours ajouter_

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