How to create a pygame surface from a numpy array of float32?

convert data range to range[0-255], the data size must be M x N or M x N x 3

pygame.init()
display = pygame.display.set_mode((350, 350))
x = np.arange(0, 300)
y = np.arange(0, 300)
X, Y = np.meshgrid(x, y)
Z = X + Y
Z = 255*Z/Z.max()
surf = pygame.surfarray.make_surface(Z)

running = True

while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
    display.blit(surf, (0, 0))
    pygame.display.update()
pygame.quit()

enter image description here

If you want in grayscale:

import pygame
import numpy as np


def gray(im):
    im = 255 * (im / im.max())
    w, h = im.shape
    ret = np.empty((w, h, 3), dtype=np.uint8)
    ret[:, :, 2] = ret[:, :, 1] = ret[:, :, 0] = im
    return ret

pygame.init()
display = pygame.display.set_mode((350, 350))
x = np.arange(0, 300)
y = np.arange(0, 300)
X, Y = np.meshgrid(x, y)
Z = X + Y
Z = 255 * Z / Z.max()
Z = gray(Z)
surf = pygame.surfarray.make_surface(Z)

running = True

while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
    display.blit(surf, (0, 0))
    pygame.display.update()
pygame.quit()

enter image description here


Let's say for simplicity you have only values [0,1], otherwize it is better to truncate values to a fixed range, so that by some float value, e.g. 0,5 you have always same color output.
I'll take simple horizontal gradient as input example:

W = 300
H = 200
# generate simple gradient in float
F0 = numpy.linspace(0, 1, num = W)
F = numpy.tile(F0, (H, 1))

Now there are couple of ways to show it up. I would probaby in this case show it on 8 bit surface. In this case you'll need this to define the color palette in Pygame format:

def make_palette (C1, C2):
    palR = numpy.linspace(C1[0], C2[0], num = 256, dtype = "uint8")
    palG = numpy.linspace(C1[1], C2[1], num = 256, dtype = "uint8")
    palB = numpy.linspace(C1[2], C2[2], num = 256, dtype = "uint8")
    return zip(palR,palG,palB)

And this to copy data from array to surface:

def put_arr(Dest, Src):                 
    buf = Dest.get_buffer()
    buf.write(Src.tostring(), 0)

Now in the beginning of the program you initialize the surface of the same size as you array and apply palette:

I_surf = pygame.Surface((W, H), 0, 8)           # Pygame surface 
C1 = (0,0,250)
C2 = (250,0,0)
palRGB = make_palette (C1, C2)
I_surf.set_palette(palRGB)

And in the main loop you have something like:

I = numpy.rint( F*255 ).astype("uint8")
put_arr(I_surf, I)
...
DISPLAY.blit(I_surf, (100, 100))

Note the type of array and surface, they both must be 8 bit in this case.
If all works, you must see this in the window:

enter image description here


This is a supplement to eyllanesc's answer, since I found that answer's examples helpful, but they can be optimised slightly if you want to update the image.

The main optimisations are to use an 8-bit array, and to use surfarray.blit_array instead of blit and display.flip instead of display.update. (The latter two require the display and the array to be the same size.)

The difference is not huge, however - the code below gets me about 15 fps without the optimisations and 20-21 fps with them. An OpenGL based method would be much faster, but I don't know a convenient library for that.

import pygame
import numpy as np
from pygame import surfarray

from timeit import default_timer as timer

pygame.init()

w = 1000
h = 800

display = pygame.display.set_mode((w, h))

img = np.zeros((w,h,3),dtype=np.uint8)


init_time = timer()
frames_displayed = 0

running = True
while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False

    # updating the image
    for i in range(100):
        img[np.random.randint(w),np.random.randint(h)] = np.random.randint(255,size=3,dtype=np.uint8)

    surfarray.blit_array(display, img)
    pygame.display.flip()

    frames_displayed+=1


print("average frame rate:", frames_displayed/(timer()-init_time), "fps")

pygame.quit()