Expyriment: PixelArray and numpy Demo

expy_numpy_demo.py

import pygame
import pygame.surfarray as surfarray
import numpy
from expyriment import control
from expyriment.stimuli import Canvas

class PGSurface(Canvas):
    """PyGame Surface: Expyriment Stimulus for direct Pygame operations and
    PixelArrays

    In contrast to other Expyriment stimuli the class does not generate temporary
    surfaces.
    """

    def __init__(self, size, position=None, colour=None):
        Canvas.__init__(self, size, position, colour)
        self._px_array = None

    @property
    def surface(self):
        """todo"""
        if not self.has_surface:
            ok = self._set_surface(self._get_surface())  # create surface
            if not ok:
                raise RuntimeError(Visual._compression_exception_message.format(
                    "surface"))
        return self._surface

    @property
    def pixel_array(self):
        """todo"""
        if self._px_array is None:
            self._px_array = pygame.PixelArray(self.surface)
        return self._px_array

    @pixel_array.setter
    def pixel_array(self, value):
        if self._px_array is None:
            self._px_array = pygame.PixelArray(self.surface)
        self._px_array = value

    def get_numpy_array(self):
        self.unlock_pixel_array()
        return surfarray.pixels3d(self.surface)

    def set_numpy_array(self, array):
        self.unlock_pixel_array()
        surfarray.blit_array(self.surface, array)

    def unlock_pixel_array(self):
        """todo"""
        self._px_array = None

    def preload(self, inhibit_ogl_compress=False):
        self.unlock_pixel_array()
        return Canvas.preload(self, inhibit_ogl_compress)

    def compress(self):
        self.unlock_pixel_array()
        return Canvas.compress(self)

    def decompress(self):
        self.unlock_pixel_array()
        return Canvas.decompress(self)

    def plot(self, stimulus):
        self.unlock_pixel_array()
        return Canvas.plot(self, stimulus)

    def clear_surface(self):
        self.unlock_pixel_array()
        return Canvas.clear_surface(self)

    def copy(self):
        self.unlock_pixel_array()
        return Canvas.copy(self)

    def unload(self, keep_surface=False):
        if not keep_surface:
            self.unlock_pixel_array()
        return Canvas.unload(self, keep_surface)

    def rotate(self, degree):
        self.unlock_pixel_array()
        return Canvas.rotate(self, degree)

    def scale(self, factors):
        self.unlock_pixel_array()
        return Canvas.scale(self, factors)

    # expyriment 0.8.0
    # def scale_to_fullscreen(self, keep_aspect_ratio=True):
    # self.unlock_pixel_array()
    # return Canvas.scale_to_fullscreen(self, keep_aspect_ratio)

    def flip(self, booleans):
        self.unlock_pixel_array()
        return Canvas.flip(self, booleans)

    def blur(self, level):
        self.unlock_pixel_array()
        return Canvas.blur(self, level)

    def scramble(self, grain_size):
        self.unlock_pixel_array()
        return Canvas.scramble(self, grain_size)

    def add_noise(self, grain_size, percentage, colour):
        self.unlock_pixel_array()
        return Canvas.add_noise(self, grain_size, percentage, colour)


control.set_develop_mode(True)
exp = control.initialize()
control.start(skip_ready_screen=True)


pgsurface = PGSurface(size=(200, 200))
pgsurface.pixel_array[:, :] = (255, 0, 0) # set all to red

#make numpy array (3d for rgb)
nparr = pgsurface.get_numpy_array()
# change values in 
nparr[:, 30:32, :] = 0
nparr[:, 30:32, 1] = 255
# convert to surface
pgsurface.set_numpy_array(nparr)

pgsurface.present()


exp.keyboard.wait()
control.end()

PGSurface could be included as Expyriment Stimulus?