Source

pygame / lib / _numpysurfarray.py

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
##    pygame - Python Game Library
##    Copyright (C) 2007 Marcus von Appen
##
##    This library is free software; you can redistribute it and/or
##    modify it under the terms of the GNU Library General Public
##    License as published by the Free Software Foundation; either
##    version 2 of the License, or (at your option) any later version.
##
##    This library is distributed in the hope that it will be useful,
##    but WITHOUT ANY WARRANTY; without even the implied warranty of
##    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
##    Library General Public License for more details.
##
##    You should have received a copy of the GNU Library General Public
##    License along with this library; if not, write to the Free
##    Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
##
##    Marcus von Appen
##    mva@sysfault.org

"""pygame module for accessing surface pixel data using numpy

Functions to convert pixel data between pygame Surfaces and Numpy
arrays. This module will only be available when pygame can use the
external Numpy package.

Note, that numpyarray is an optional module. It requires that Numpy is
installed to be used. If not installed, an exception will be raised when
it is used. eg. ImportError: no module named numpy

Every pixel is stored as a single integer value to represent the red,
green, and blue colors. The 8bit images use a value that looks into a
colormap. Pixels with higher depth use a bit packing process to place
three or four values into a single number.

The Numpy arrays are indexed by the X axis first, followed by the Y
axis. Arrays that treat the pixels as a single integer are referred to
as 2D arrays. This module can also separate the red, green, and blue
color values into separate indices. These types of arrays are referred
to as 3D arrays, and the last index is 0 for red, 1 for green, and 2 for
blue.

In contrast to Numeric Numpy does use unsigned 16bit integers, images
with 16bit data will be treated as unsigned integers.
"""

import pygame
import numpy
import re

def array2d (surface):
    """pygame.numpyarray.array2d (Surface): return array

    copy pixels into a 2d array

    Copy the pixels from a Surface into a 2D array. The bit depth of the
    surface will control the size of the integer values, and will work
    for any type of pixel format.

    This function will temporarily lock the Surface as pixels are copied
    (see the Surface.lock - lock the Surface memory for pixel access
    method).
    """
    bpp = surface.get_bytesize ()
    if bpp <= 0 or bpp > 4:
        raise ValueError, "unsupported bit depth for 2D array"

    # Taken from Alex Holkner's pygame-ctypes package. Thanks a lot.
    data = surface.get_buffer ().raw
        
    # Remove extra pitch from each row.
    width = surface.get_width ()
    pitchdiff = surface.get_pitch () - width * bpp
    if pitchdiff > 0:
        pattern = re.compile ('(%s)%s' % ('.' * width * bpp, '.' * pitchdiff),
                              flags=re.DOTALL)
        data = ''.join (pattern.findall (data))

    if bpp == 3:
        # Pad each triplet of bytes with another zero
        pattern = re.compile ('...', flags=re.DOTALL)
        data = '\0'.join (pattern.findall (data))
        if pygame.get_sdl_byteorder () == pygame.LIL_ENDIAN:
            data += '\0'
        else:
            data = '\0' + data
        bpp = 4

    typecode = (numpy.uint8, numpy.uint16, None, numpy.uint32)[bpp - 1]
    array = numpy.fromstring (data, typecode)
    array.shape = (surface.get_height (), width)
    array = numpy.transpose (array)
    return array
    

def pixels2d (surface):
    """pygame.numpyarray.pixels2d (Surface): return array

    reference pixels into a 2d array
    
    Create a new 2D array that directly references the pixel values in a
    Surface. Any changes to the array will affect the pixels in the
    Surface. This is a fast operation since no data is copied.

    Pixels from a 24-bit Surface cannot be referenced, but all other
    Surface bit depths can.

    The Surface this references will remain locked for the lifetime of
    the array (see the Surface.lock - lock the Surface memory for pixel
    access method).
    """
    bpp = surface.get_bytesize ()
    if bpp == 3 or bpp < 1 or bpp > 4:
        raise ValueError, "unsupported bit depth for 2D reference array"

    typecode = (numpy.uint8, numpy.uint16, None, numpy.uint32)[bpp - 1]
    array = numpy.frombuffer (surface.get_buffer (), typecode)
    array.shape = surface.get_height (), surface.get_pitch () / bpp

    # Padding correction for certain depth due to padding bytes.
    array = array[:,:surface.get_width ()]
    array = numpy.transpose (array)
    return array

def array3d (surface):
    """pygame.numpyarray.array3d (Surface): return array

    copy pixels into a 3d array

    Copy the pixels from a Surface into a 3D array. The bit depth of the
    surface will control the size of the integer values, and will work
    for any type of pixel format.

    This function will temporarily lock the Surface as pixels are copied
    (see the Surface.lock - lock the Surface memory for pixel access
    method).
    """
    bpp = surface.get_bytesize ()
    array = array2d (surface)

    # Taken from from Alex Holkner's pygame-ctypes package. Thanks a
    # lot.
    if bpp == 1:
        palette = surface.get_palette ()
        # Resolve the correct values using the color palette
        pal_r = numpy.array ([c[0] for c in palette])
        pal_g = numpy.array ([c[1] for c in palette])
        pal_b = numpy.array ([c[2] for c in palette])
        planes = [numpy.choose (array, pal_r),
                  numpy.choose (array, pal_g),
                  numpy.choose (array, pal_b)]
        array = numpy.array (planes)
        array = numpy.transpose (array, (1, 2, 0))
        return array
    else:
        masks = surface.get_masks ()
        shifts = surface.get_shifts ()
        losses = surface.get_losses ()
        planes = [((array & masks[0]) >> shifts[0]) << losses[0],
                  ((array & masks[1]) >> shifts[1]) << losses[1],
                  ((array & masks[2]) >> shifts[2]) << losses[2]]
        array = numpy.array (planes)
        return numpy.transpose (array, (1, 2, 0))

def pixels3d (surface):
    """pygame.numpyarray.pixels3d (Surface): return array

    reference pixels into a 3d array

    Create a new 3D array that directly references the pixel values in a
    Surface. Any changes to the array will affect the pixels in the
    Surface. This is a fast operation since no data is copied.

    This will only work on Surfaces that have 24-bit or 32-bit
    formats. Lower pixel formats cannot be referenced.

    The Surface this references will remain locked for the lifetime of
    the array (see the Surface.lock - lock the Surface memory for pixel
    access method).
    """
    bpp = surface.get_bytesize ()
    if bpp < 3 or bpp > 4:
        raise ValueError, "unsupported bit depth for 3D reference array"

    start = 0
    step = 0
    end = 0

    # Check for RGB or BGR surface.
    shifts = surface.get_shifts ()
    if shifts[0] == 16 and shifts[1] == 8 and shifts[2] == 0:
        # RGB 
        end = None
        if pygame.get_sdl_byteorder () == pygame.LIL_ENDIAN:
            start = 2
            step = -1
        else:
            start = 0
            step = 1
    else:
        # BGR
        end = 3
        if pygame.get_sdl_byteorder () == pygame.LIL_ENDIAN:
            start = 0
            step = 1
        else:
            start = 2
            step = -1

    array = numpy.frombuffer (surface.get_buffer (), numpy.uint8)
    array.shape = surface.get_height (), surface.get_pitch ()
    array = array[:,:surface.get_width () * bpp]
    array.shape  = surface.get_width (), surface.get_height (), bpp
    array = array[:,:,start:end:step]
    return array

def array_alpha (surface):
    """pygame.numpyarray.array_alpha (Surface): return array

    copy pixel alphas into a 2d array

    Copy the pixel alpha values (degree of transparency) from a Surface
    into a 2D array. This will work for any type of Surface
    format. Surfaces without a pixel alpha will return an array with all
    opaque values.

    This function will temporarily lock the Surface as pixels are copied
    (see the Surface.lock - lock the Surface memory for pixel access
    method).
    """
    if surface.get_bytesize () == 1 or not surface.get_alpha ():
        # 1 bpp surfaces and surfaces without alpha are always fully
        # opaque.
        array = numpy.empty (surface.get_width () * surface.get_height (),
                             numpy.uint8)
        array.fill (0xff)
        array.shape = surface.get_width (), surface.get_height ()
        return array

    array = array2d (surface)
    # Those shifts match the results from the old numeric system
    # exactly.
    array = array >> surface.get_shifts ()[3] << surface.get_losses ()[3]
    array = array.astype (numpy.uint8)
    return array

def pixels_alpha (surface):
    """pygame.numpyarray.pixels_alpha (Surface): return array

    reference pixel alphas into a 2d array

    Create a new 2D array that directly references the alpha values
    (degree of transparency) in a Surface. Any changes to the array will
    affect the pixels in the Surface. This is a fast operation since no
    data is copied.

    This can only work on 32-bit Surfaces with a per-pixel alpha value.

    The Surface this array references will remain locked for the
    lifetime of the array.
    """
    if surface.get_bytesize () != 4:
        raise ValueError, "unsupported bit depth for alpha reference array"

    # ARGB surface.
    start = 0
    
    if surface.get_shifts ()[3] == 24:
        # RGBA surface.
        start = 3

    width, height = surface.get_width (), surface.get_height ()
    array = numpy.frombuffer (surface.get_buffer (), numpy.uint8)
    array.shape = height, surface.get_pitch ()
    array = array[:, start::4]
    array = array[:, :width * 4]
    return numpy.transpose (array)

def array_colorkey (surface):
    """pygame.numpyarray.array_colorkey (Surface): return array

    copy the colorkey values into a 2d array

    Create a new array with the colorkey transparency value from each
    pixel. If the pixel matches the colorkey it will be fully
    tranparent; otherwise it will be fully opaque.

    This will work on any type of Surface format. If the image has no
    colorkey a solid opaque array will be returned.

    This function will temporarily lock the Surface as pixels are
    copied.
    """
    colorkey = surface.get_colorkey ()
    if colorkey == None:
        # No colorkey, return a solid opaque array.
        array = numpy.empty (surface.get_width () * surface.get_height (),
                             numpy.uint8)
        array.fill (0xff)
        array.shape = surface.get_width (), surface.get_height ()
        return array

    # Taken from from Alex Holkner's pygame-ctypes package. Thanks a
    # lot.
    array = array2d (surface)
    # Check each pixel value for the colorkey and mark it as opaque or
    # transparent as needed.
    val = surface.map_rgb (colorkey)
    array = numpy.choose (numpy.equal (array, val),
                          (numpy.uint8 (0xff), numpy.uint8 (0)))
    array.shape = surface.get_width (), surface.get_height ()
    return array

def make_surface (array):
    """pygame.numpyarray.make_surface (array): return Surface

    copy an array to a new surface

    Create a new Surface that best resembles the data and format on the
    array. The array can be 2D or 3D with any sized integer values.
    """ 
    # Taken from from Alex Holkner's pygame-ctypes package. Thanks a
    # lot.
    bpp = 0
    r = g = b = 0
    shape = array.shape
    if len (shape) == 2:
        # 2D array
        bpp = 8
        
    elif len (shape) == 3 and shape[2] == 3:
        bpp = 32
        r = 0xff << 16
        g = 0xff << 8
        b = 0xff
    else:
        raise ValueError, "must be a valid 2d or 3d array"

    surface = pygame.Surface ((shape[0], shape[1]), 0, bpp, (r, g, b, 0))
    blit_array (surface, array)
    return surface

def blit_array (surface, array):
    """pygame.numpyarray.blit_array (Surface, array): return None

    blit directly from a array values

    Directly copy values from an array into a Surface. This is faster
    than converting the array into a Surface and blitting. The array
    must be the same dimensions as the Surface and will completely
    replace all pixel values.

    This function will temporarily lock the Surface as the new values
    are copied.
    """
    bpp = surface.get_bytesize ()
    if bpp <= 0 or bpp > 4:
        raise ValueError, "unsupported bit depth for surface"
    
    shape = array.shape
    width = surface.get_width ()

    typecode = (numpy.uint8, numpy.uint16, None, numpy.uint32)[bpp - 1]
    array = array.astype (typecode)

    # Taken from from Alex Holkner's pygame-ctypes package. Thanks a
    # lot.
    if len(shape) == 3 and shape[2] == 3:
        array = numpy.transpose (array, (1, 0, 2))
        shifts = surface.get_shifts ()
        losses = surface.get_losses ()
        array = (array[:,:,::3] >> losses[0] << shifts[0]) | \
                (array[:,:,1::3] >> losses[1] << shifts[1]) | \
                (array[:,:,2::3] >> losses[2] << shifts[2])
    elif len (shape) == 2:
        array = numpy.transpose (array)
    else:
        raise ValueError, "must be a valid 2d or 3d array"

    if width != shape[0] or surface.get_height () != shape[1]:
        raise ValueError, "array must match the surface dimensions"

    itemsize = array.itemsize
    data = array.tostring ()
    
    if itemsize > bpp:
        # Trim bytes from each element, keep least significant byte(s)
        pattern = '%s(%s)' % ('.' * (itemsize - bpp), '.' * bpp)
        if pygame.get_sdl_byteorder () == pygame.LIL_ENDIAN:
            pattern = '(%s)%s' % ('.' * bpp, '.' * (itemsize - bpp))
        data = ''.join (re.compile (pattern, flags=re.DOTALL).findall (data))
    elif itemsize < bpp:
        # Add pad bytes to each element, at most significant end
        pad = '\0' * (bpp - itemsize)
        pixels = re.compile ('.' * itemsize, flags=re.DOTALL).findall (data)
        data = pad.join (pixels)
        if pygame.get_sdl_byteorder () == pygame.LIL_ENDIAN:
            data = data + pad
        else:
            data = pad + data

    # Add zeros pad for pitch correction
    pitchdiff = surface.get_pitch () - width * bpp
    if pitchdiff > 0:
        pad = '\0' * pitchdiff
        rows = re.compile ('.' * width * bpp, flags=re.DOTALL).findall (data)
        data = pad.join (rows) + pad

    surface.get_buffer ().write (data, 0)
    
def map_array (surface, array):
    """pygame.numpyarray.map_array (Surface, array3d): return array2d

    map a 3d array into a 2d array

    Convert a 3D array into a 2D array. This will use the given Surface
    format to control the conversion. Palette surface formats are not
    supported.

    Note: arrays do not need to be 3D, as long as the minor axis has
    three elements giving the component colours, any array shape can be
    used (for example, a single colour can be mapped, or an array of
    colours).
    """
    # Taken from from Alex Holkner's pygame-ctypes package. Thanks a
    # lot.
    bpp = surface.get_bytesize ()
    if bpp <= 0 or bpp > 4:
        raise ValueError, "unsupported bit depth for surface array"

    shape = array.shape
    if shape[-1] != 3:
        raise ValueError, "array must be a 3d array of 3-value color data"

    shifts = surface.get_shifts ()
    losses = surface.get_losses ()
    array = (array[:,:,::3] >> losses[0] << shifts[0]) | \
            (array[:,:,1::3] >> losses[1] << shifts[1]) | \
            (array[:,:,2::3] >> losses[2] << shifts[2])
    return array
Tip: Filter by directory path e.g. /media app.js to search for public/media/app.js.
Tip: Use camelCasing e.g. ProjME to search for ProjectModifiedEvent.java.
Tip: Filter by extension type e.g. /repo .js to search for all .js files in the /repo directory.
Tip: Separate your search with spaces e.g. /ssh pom.xml to search for src/ssh/pom.xml.
Tip: Use ↑ and ↓ arrow keys to navigate and return to view the file.
Tip: You can also navigate files with Ctrl+j (next) and Ctrl+k (previous) and view the file with Ctrl+o.
Tip: You can also navigate files with Alt+j (next) and Alt+k (previous) and view the file with Alt+o.