_numpysurfarray.py 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356
  1. ## pygame - Python Game Library
  2. ## Copyright (C) 2007 Marcus von Appen
  3. ##
  4. ## This library is free software; you can redistribute it and/or
  5. ## modify it under the terms of the GNU Library General Public
  6. ## License as published by the Free Software Foundation; either
  7. ## version 2 of the License, or (at your option) any later version.
  8. ##
  9. ## This library is distributed in the hope that it will be useful,
  10. ## but WITHOUT ANY WARRANTY; without even the implied warranty of
  11. ## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
  12. ## Library General Public License for more details.
  13. ##
  14. ## You should have received a copy of the GNU Library General Public
  15. ## License along with this library; if not, write to the Free
  16. ## Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
  17. ##
  18. ## Marcus von Appen
  19. ## mva@sysfault.org
  20. """pygame module for accessing surface pixel data using numpy
  21. Functions to convert pixel data between pygame Surfaces and Numpy
  22. arrays. This module will only be available when pygame can use the
  23. external Numpy package.
  24. Note, that numpyarray is an optional module. It requires that Numpy is
  25. installed to be used. If not installed, an exception will be raised when
  26. it is used. eg. ImportError: no module named numpy
  27. Every pixel is stored as a single integer value to represent the red,
  28. green, and blue colors. The 8bit images use a value that looks into a
  29. colormap. Pixels with higher depth use a bit packing process to place
  30. three or four values into a single number.
  31. The Numpy arrays are indexed by the X axis first, followed by the Y
  32. axis. Arrays that treat the pixels as a single integer are referred to
  33. as 2D arrays. This module can also separate the red, green, and blue
  34. color values into separate indices. These types of arrays are referred
  35. to as 3D arrays, and the last index is 0 for red, 1 for green, and 2 for
  36. blue.
  37. In contrast to Numeric Numpy does use unsigned 16bit integers, images
  38. with 16bit data will be treated as unsigned integers.
  39. """
  40. import pygame
  41. from pygame.compat import bytes_
  42. from pygame.pixelcopy import array_to_surface, surface_to_array, \
  43. map_array as pix_map_array, make_surface as pix_make_surface
  44. import numpy
  45. from numpy import array as numpy_array, empty as numpy_empty, \
  46. around as numpy_around, uint32 as numpy_uint32, \
  47. ndarray as numpy_ndarray
  48. #float96 not available on all numpy versions.
  49. numpy_floats = []
  50. for type_name in "float float32 float64 float96".split():
  51. if hasattr(numpy, type_name):
  52. numpy_floats.append(getattr(numpy, type_name))
  53. # Pixel sizes corresponding to NumPy supported integer sizes, and therefore
  54. # permissible for 2D reference arrays.
  55. _pixel2d_bitdepths = set([8, 16, 32])
  56. def blit_array (surface, array):
  57. """pygame.surfarray.blit_array(Surface, array): return None
  58. Blit directly from a array values.
  59. Directly copy values from an array into a Surface. This is faster than
  60. converting the array into a Surface and blitting. The array must be the
  61. same dimensions as the Surface and will completely replace all pixel
  62. values. Only integer, ascii character and record arrays are accepted.
  63. This function will temporarily lock the Surface as the new values are
  64. copied.
  65. """
  66. if isinstance(array, numpy_ndarray) and array.dtype in numpy_floats:
  67. array = array.round(0).astype(numpy_uint32)
  68. return array_to_surface(surface, array)
  69. def make_surface(array):
  70. """pygame.surfarray.make_surface (array): return Surface
  71. Copy an array to a new surface.
  72. Create a new Surface that best resembles the data and format on the
  73. array. The array can be 2D or 3D with any sized integer values.
  74. """
  75. if isinstance(array, numpy_ndarray) and array.dtype in numpy_floats:
  76. array = array.round(0).astype(numpy_uint32)
  77. return pix_make_surface (array)
  78. def array2d(surface):
  79. """pygame.numpyarray.array2d(Surface): return array
  80. copy pixels into a 2d array
  81. Copy the pixels from a Surface into a 2D array. The bit depth of the
  82. surface will control the size of the integer values, and will work
  83. for any type of pixel format.
  84. This function will temporarily lock the Surface as pixels are copied
  85. (see the Surface.lock - lock the Surface memory for pixel access
  86. method).
  87. """
  88. bpp = surface.get_bytesize()
  89. try:
  90. dtype = (numpy.uint8, numpy.uint16, numpy.int32, numpy.int32)[bpp - 1]
  91. except IndexError:
  92. raise ValueError("unsupported bit depth %i for 2D array" % (bpp * 8,))
  93. size = surface.get_size()
  94. array = numpy.empty(size, dtype)
  95. surface_to_array(array, surface)
  96. return array
  97. def pixels2d(surface):
  98. """pygame.numpyarray.pixels2d(Surface): return array
  99. reference pixels into a 2d array
  100. Create a new 2D array that directly references the pixel values in a
  101. Surface. Any changes to the array will affect the pixels in the
  102. Surface. This is a fast operation since no data is copied.
  103. Pixels from a 24-bit Surface cannot be referenced, but all other
  104. Surface bit depths can.
  105. The Surface this references will remain locked for the lifetime of
  106. the array (see the Surface.lock - lock the Surface memory for pixel
  107. access method).
  108. """
  109. if (surface.get_bitsize() not in _pixel2d_bitdepths):
  110. raise ValueError("unsupport bit depth for 2D reference array")
  111. try:
  112. return numpy_array(surface.get_view('2'), copy=False)
  113. except (ValueError, TypeError):
  114. raise ValueError("bit depth %i unsupported for 2D reference array" %
  115. (surface.get_bitsize(),))
  116. def array3d(surface):
  117. """pygame.numpyarray.array3d(Surface): return array
  118. copy pixels into a 3d array
  119. Copy the pixels from a Surface into a 3D array. The bit depth of the
  120. surface will control the size of the integer values, and will work
  121. for any type of pixel format.
  122. This function will temporarily lock the Surface as pixels are copied
  123. (see the Surface.lock - lock the Surface memory for pixel access
  124. method).
  125. """
  126. w, h = surface.get_size()
  127. array = numpy.empty((w, h, 3), numpy.uint8)
  128. surface_to_array(array, surface)
  129. return array
  130. def pixels3d (surface):
  131. """pygame.numpyarray.pixels3d(Surface): return array
  132. reference pixels into a 3d array
  133. Create a new 3D array that directly references the pixel values in a
  134. Surface. Any changes to the array will affect the pixels in the
  135. Surface. This is a fast operation since no data is copied.
  136. This will only work on Surfaces that have 24-bit or 32-bit
  137. formats. Lower pixel formats cannot be referenced.
  138. The Surface this references will remain locked for the lifetime of
  139. the array (see the Surface.lock - lock the Surface memory for pixel
  140. access method).
  141. """
  142. return numpy_array(surface.get_view('3'), copy=False)
  143. def array_alpha(surface):
  144. """pygame.numpyarray.array_alpha(Surface): return array
  145. copy pixel alphas into a 2d array
  146. Copy the pixel alpha values (degree of transparency) from a Surface
  147. into a 2D array. This will work for any type of Surface
  148. format. Surfaces without a pixel alpha will return an array with all
  149. opaque values.
  150. This function will temporarily lock the Surface as pixels are copied
  151. (see the Surface.lock - lock the Surface memory for pixel access
  152. method).
  153. """
  154. size = surface.get_size()
  155. array = numpy.empty(size, numpy.uint8)
  156. surface_to_array(array, surface, 'A')
  157. return array
  158. def pixels_alpha(surface):
  159. """pygame.numpyarray.pixels_alpha(Surface): return array
  160. reference pixel alphas into a 2d array
  161. Create a new 2D array that directly references the alpha values
  162. (degree of transparency) in a Surface. Any changes to the array will
  163. affect the pixels in the Surface. This is a fast operation since no
  164. data is copied.
  165. This can only work on 32-bit Surfaces with a per-pixel alpha value.
  166. The Surface this array references will remain locked for the
  167. lifetime of the array.
  168. """
  169. return numpy.array(surface.get_view('A'), copy=False)
  170. def pixels_red(surface):
  171. """pygame.surfarray.pixels_red(Surface): return array
  172. Reference pixel red into a 2d array.
  173. Create a new 2D array that directly references the red values
  174. in a Surface. Any changes to the array will affect the pixels
  175. in the Surface. This is a fast operation since no data is copied.
  176. This can only work on 24-bit or 32-bit Surfaces.
  177. The Surface this array references will remain locked for the
  178. lifetime of the array.
  179. """
  180. return numpy.array(surface.get_view('R'), copy=False)
  181. def array_red(surface):
  182. """pygame.numpyarray.array_red(Surface): return array
  183. copy pixel red into a 2d array
  184. Copy the pixel red values from a Surface into a 2D array. This will work
  185. for any type of Surface format.
  186. This function will temporarily lock the Surface as pixels are copied
  187. (see the Surface.lock - lock the Surface memory for pixel access
  188. method).
  189. """
  190. size = surface.get_size()
  191. array = numpy.empty(size, numpy.uint8)
  192. surface_to_array(array, surface, 'R')
  193. return array
  194. def pixels_green(surface):
  195. """pygame.surfarray.pixels_green(Surface): return array
  196. Reference pixel green into a 2d array.
  197. Create a new 2D array that directly references the green values
  198. in a Surface. Any changes to the array will affect the pixels
  199. in the Surface. This is a fast operation since no data is copied.
  200. This can only work on 24-bit or 32-bit Surfaces.
  201. The Surface this array references will remain locked for the
  202. lifetime of the array.
  203. """
  204. return numpy.array(surface.get_view('G'), copy=False)
  205. def array_green(surface):
  206. """pygame.numpyarray.array_green(Surface): return array
  207. copy pixel green into a 2d array
  208. Copy the pixel green values from a Surface into a 2D array. This will work
  209. for any type of Surface format.
  210. This function will temporarily lock the Surface as pixels are copied
  211. (see the Surface.lock - lock the Surface memory for pixel access
  212. method).
  213. """
  214. size = surface.get_size()
  215. array = numpy.empty(size, numpy.uint8)
  216. surface_to_array(array, surface, 'G')
  217. return array
  218. def pixels_blue (surface):
  219. """pygame.surfarray.pixels_blue(Surface): return array
  220. Reference pixel blue into a 2d array.
  221. Create a new 2D array that directly references the blue values
  222. in a Surface. Any changes to the array will affect the pixels
  223. in the Surface. This is a fast operation since no data is copied.
  224. This can only work on 24-bit or 32-bit Surfaces.
  225. The Surface this array references will remain locked for the
  226. lifetime of the array.
  227. """
  228. return numpy.array(surface.get_view('B'), copy=False)
  229. def array_blue(surface):
  230. """pygame.numpyarray.array_blue(Surface): return array
  231. copy pixel blue into a 2d array
  232. Copy the pixel blue values from a Surface into a 2D array. This will work
  233. for any type of Surface format.
  234. This function will temporarily lock the Surface as pixels are copied
  235. (see the Surface.lock - lock the Surface memory for pixel access
  236. method).
  237. """
  238. size = surface.get_size()
  239. array = numpy.empty(size, numpy.uint8)
  240. surface_to_array(array, surface, 'B')
  241. return array
  242. def array_colorkey(surface):
  243. """pygame.numpyarray.array_colorkey(Surface): return array
  244. copy the colorkey values into a 2d array
  245. Create a new array with the colorkey transparency value from each
  246. pixel. If the pixel matches the colorkey it will be fully
  247. tranparent; otherwise it will be fully opaque.
  248. This will work on any type of Surface format. If the image has no
  249. colorkey a solid opaque array will be returned.
  250. This function will temporarily lock the Surface as pixels are
  251. copied.
  252. """
  253. size = surface.get_size()
  254. array = numpy.empty(size, numpy.uint8)
  255. surface_to_array(array, surface, 'C')
  256. return array
  257. def map_array(surface, array):
  258. """pygame.numpyarray.map_array(Surface, array3d): return array2d
  259. map a 3d array into a 2d array
  260. Convert a 3D array into a 2D array. This will use the given Surface
  261. format to control the conversion.
  262. Note: arrays do not need to be 3D, as long as the minor axis has
  263. three elements giving the component colours, any array shape can be
  264. used (for example, a single colour can be mapped, or an array of
  265. colours). The array shape is limited to eleven dimensions maximum,
  266. including the three element minor axis.
  267. """
  268. if array.ndim == 0:
  269. raise ValueError("array must have at least 1 dimension")
  270. shape = array.shape
  271. if shape[-1] != 3:
  272. raise ValueError("array must be a 3d array of 3-value color data")
  273. target = numpy_empty(shape[:-1], numpy.int32)
  274. pix_map_array(target, array, surface)
  275. return target