Source code for sdl2.ext.pixelaccess

"""Pixel-wise access routines."""
import ctypes
from .compat import UnsupportedError, experimental
from .array import MemoryView
from ..surface import SDL_MUSTLOCK, SDL_LockSurface, SDL_UnlockSurface, \
    SDL_Surface
from ..stdinc import Uint8
from .sprite import SoftwareSprite
from .draw import _get_target_surface, prepare_color


__all__ = ["PixelView", "pixels2d", "pixels3d"]


[docs]class PixelView(MemoryView): """2D memory view for Sprite and SDL_Surface pixel access. The PixelView uses a y/x-layout. Accessing view[N] will operate on the Nth row of the underlying surface. To access a specific column within that row, view[N][C] has to be used. NOTE: The PixelView is implemented on top of the MemoryView class. As such it makes heavy use of recursion to access rows and columns and can be considered as slow in contrast to optimised ndim-array solutions such as numpy. """ def __init__(self, source): """Creates a new PixelView from a Sprite or SDL_Surface. If necessary, the surface will be locked for accessing its pixel data. The lock will be removed once the PixelView is garbage-collected or deleted. """ if isinstance(source, SoftwareSprite): self._surface = source.surface # keep a reference, so the Sprite's not GC'd self._sprite = source elif isinstance(source, SDL_Surface): self._surface = source elif "SDL_Surface" in str(type(source)): self._surface = source.contents else: raise TypeError("source must be a Sprite or SDL_Surface") if SDL_MUSTLOCK(self._surface): SDL_LockSurface(self._surface) pxbuf = ctypes.cast(self._surface.pixels, ctypes.POINTER(Uint8)) itemsize = self._surface.format.contents.BytesPerPixel strides = (self._surface.h, self._surface.w) srcsize = self._surface.h * self._surface.pitch super(PixelView, self).__init__(pxbuf, itemsize, strides, getfunc=self._getitem, setfunc=self._setitem, srcsize=srcsize) def _getitem(self, start, end): if self.itemsize == 1: # byte-wise access return self.source[start:end] # move the pointer to the correct location src = ctypes.byref(self.source.contents, start) casttype = ctypes.c_ubyte if self.itemsize == 2: casttype = ctypes.c_ushort elif self.itemsize == 3: # TODO raise NotImplementedError("unsupported bpp") elif self.itemsize == 4: casttype = ctypes.c_uint return ctypes.cast(src, ctypes.POINTER(casttype)).contents.value def _setitem(self, start, end, value): target = None if self.itemsize == 1: target = ctypes.cast(self.source, ctypes.POINTER(ctypes.c_ubyte)) elif self.itemsize == 2: target = ctypes.cast(self.source, ctypes.POINTER(ctypes.c_ushort)) elif self.itemsize == 3: # TODO raise NotImplementedError("unsupported bpp") elif self.itemsize == 4: target = ctypes.cast(self.source, ctypes.POINTER(ctypes.c_uint)) value = prepare_color(value, self._surface) target[start // self.itemsize] = value def __del__(self): if self._surface: if SDL_MUSTLOCK(self._surface): SDL_UnlockSurface(self._surface)
_HASNUMPY = True try: import numpy class SurfaceArray(numpy.ndarray): """Wrapper class around numpy.ndarray. Used to keep track of the original source object for pixels2d() and pixels3d() to avoid the deletion of the source object. """ def __new__(cls, shape, dtype=float, buffer_=None, offset=0, strides=None, order=None, source=None, surface=None): sfarray = numpy.ndarray.__new__(cls, shape, dtype, buffer_, offset, strides, order) sfarray._source = source sfarray._surface = surface return sfarray def __array_finalize__(self, sfarray): if sfarray is None: return self._source = getattr(sfarray, '_source', None) self._surface = getattr(sfarray, '_surface', None) def __del__(self): if self._surface: if SDL_MUSTLOCK(self._surface): SDL_UnlockSurface(self._surface) except ImportError: _HASNUMPY = False @experimental def pixels2d(source, transpose=True): """Creates a 2D pixel array from the passed source.""" if not _HASNUMPY: raise UnsupportedError(pixels2d, "numpy module could not be loaded") psurface = _get_target_surface(source, argname="source") bpp = psurface.format.contents.BytesPerPixel if bpp < 1 or bpp > 4: raise ValueError("unsupported bpp") strides = (psurface.pitch, bpp) sz = psurface.h * psurface.pitch shape = psurface.h, psurface.w # surface.pitch // bpp dtypes = { 1: numpy.uint8, 2: numpy.uint16, 3: numpy.uint32, 4: numpy.uint32 } if SDL_MUSTLOCK(psurface): SDL_LockSurface(psurface) pxbuf = ctypes.cast(psurface.pixels, ctypes.POINTER(ctypes.c_ubyte * sz)) arr = SurfaceArray(shape, dtypes[bpp], pxbuf.contents, 0, strides, "C", source, psurface) return arr.transpose() if transpose else arr @experimental def pixels3d(source, transpose=True): """Creates a 3D pixel array from the passed source. """ if not _HASNUMPY: raise UnsupportedError(pixels3d, "numpy module could not be loaded") psurface = _get_target_surface(source, argname="source") bpp = psurface.format.contents.BytesPerPixel if bpp < 1 or bpp > 4: raise ValueError("unsupported bpp") strides = (psurface.pitch, bpp, 1) sz = psurface.h * psurface.pitch shape = psurface.h, psurface.w, bpp if SDL_MUSTLOCK(psurface): SDL_LockSurface(psurface) pxbuf = ctypes.cast(psurface.pixels, ctypes.POINTER(ctypes.c_ubyte * sz)) arr = SurfaceArray(shape, numpy.uint8, pxbuf.contents, 0, strides, "C", source, psurface) return arr.transpose(1, 0, 2) if transpose else arr