Drawing numpy arrays with pygame

While building GPU accelerated cellular automata, I needed a quick way to dump a numpy array of pixels onto the screen. This is pretty simple to do with pygame’s surfarray API. You run this in the main drawing loop:

# Z is a numpy array
surf = pygame.surfarray.make_surface(Z)
self.display.blit(surf, (0, 0))

pygame.display.update()

Here’s a class that wraps this:

# viewer.py
# (c) 2020 Karthik Karanth, MIT License

import pygame

class Viewer:
    def __init__(self, update_func, display_size):
        self.update_func = update_func
        pygame.init()
        self.display = pygame.display.set_mode(display_size)
    
    def set_title(self, title):
        pygame.display.set_caption(title)
    
    def start(self):
        running = True
        while running:
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    running = False

            Z = self.update_func()
            surf = pygame.surfarray.make_surface(Z)
            self.display.blit(surf, (0, 0))

            pygame.display.update()

        pygame.quit()

An example usage where random data with some patterns is written to the display every frame:

#!/usr/bin/python3

# test_viewer.py
# (c) 2020 Karthik Karanth, MIT License

import numpy as np

from viewer import Viewer

def update():
    image = np.random.random((600, 600, 3)) * 255.0
    image[:,:200,0] = 255.0
    image[:,200:400,1] = 255.0
    image[:,400:,2] = 255.0
    return image.astype('uint8')

viewer = Viewer(update, (600, 600))
viewer.start()

And voila:

Obviously, this is just one way to do. If you have OpenCV as a dependency already installed, you might as well use its image display functions. But using pygame, you get all the other bells and whistles it provides(input, sound, and so on).

PUBLISHED ON FEB 28, 2020