python - Pygame and Numpy Animations

Question

Suppose I have a Numpy array myAnimation of datatype np.uint8 representing an animation (multiple frames of still 8-bit RGBA images) of shape (y,x,4,k) where y is the height, x is the width, 4 is the number of channels (red, green, blue, alpha), and the k is the number of frames in the animation.

Suppose I would like to play back the frames of this NumPy array in PyGame at a specified frame rate (say, 15 frames per second) and have the animation loop.

Is this possible to do with Pygame?

If so, how would you achieve it? Can you provide an example?

Everything I've found online involves reading in a file from disk, but it's important that I'm able to use values that are already in memory since they will change frequently while the program is running.

Answer 1

[...] shape (y,x,4,k) where y is the height, x is the width, 4 is the number of channels (red, green, blue, alpha), and the k is the number of frames [...]

Yes it is possible, with a single line of code. The following line convets a (y, x, 4, k) numpy array (data) to a list of pygme.Surface objects (surf_list):

surf_list = [pygame.image.frombuffer(d[:,:,[2, 1, 0, 3]].flatten(), (data.shape[1::-1]), 'RGBA') for d in data.transpose(3, 0, 1, 2)]

respectively

surf_list = []
for d in data.transpose(3, 0, 1, 2):
    bytes = d[:,:,[2, 1, 0, 3]].flatten()
    size = data.shape[1::-1]
    format = 'RGBA'
    surface = pygame.image.frombuffer(bytes, size, format)
    surf_list.append(surface)

Explanation:

Use numpy.traspose to bring the 3rd (frame) axis moving axis 2 to the front (see Iterating over arbitrary dimension of numpy.array) and iterate through the frames:

for d in data.transpose(3, 0, 1, 2):

Create a pygame.Surface from each frame by pygame.image.frombuffer():

surface = pygame.image.frombuffer(bytes, size, format)

pygame.image.frombuffer() has 3 arguments, bytes, size, format. bytes is the 1 dimensional byte array of pixel data. numpy.ndarray.flatten return a copy of the array collapsed into one dimension. Most likely the order of the color channels is BGRA rather than RGBA. Hence you have to swap the red and blue color channel (d[:,:,[2, 1, 0, 3]]). You can skip this if the order of the color channels is RGBA:

bytes = d[:,:,[2, 1, 0, 3]].flatten() # for BGRA

bytes = d.flatten() # for RGBA

size is a tuple with 2 elements (x, y) and specifies the size of the image. The size can be get form numpy.ndarray.shape:

size = data.shape[1::-1]

or

size = (data.shape[1], data.shape[0])

format specifies the image format and has to be 'RGBA' ('BGRA' doesn't exist):

format = 'RGBA'

---

See the minimal example, which crates a (y, x, 4, k) numpy array (data ) and converts it the a list of pygme.Surface objects (surf_list):

import pygame
import numpy as np
pygame.init()
window = pygame.display.set_mode((400, 400))
clock = pygame.time.Clock()

radius = 100
frames = 20
data = np.zeros(shape = (radius*2, radius*2, 4, frames), dtype = "uint8")
for x in range(data.shape[0]): 
    for y in range(data.shape[1]):
        px, py = x - data.shape[0]/2, y - data.shape[1]/2
        for i in range(frames):
            maxY2 = (radius*radius - px*px) * pow(abs(i-frames/2) / frames, 2)
            if px*px + py* py < radius*radius:
                if py * py < maxY2:
                    data[y, x, (0, 3), i] = 255, 255
                    if (px*px + py*py)*4 > radius*radius:
                        data[y, x, (1, 2), i] = 255, 255
                else:
                    data[y, x, (2, 3), i] = 255, 255 
                   
surf_list = [pygame.image.frombuffer(d[:,:,[2, 1, 0, 3]].flatten(), (data.shape[1::-1]), 'RGBA') for d in data.transpose(3, 0, 1, 2)]

count = 0
run = False
while not run:
    clock.tick(20)
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            run = True

    window.fill(0)
    window.blit(surf_list[count], surf_list[0].get_rect(center = window.get_rect().center))
    pygame.display.flip()
    count = (count + 1) %  len(surf_list)

pygame.quit()