1#!/usr/bin/env python
2# $URL$
3# $Rev$
4
5# Numpy example.
6# Original code created by Mel Raab, modified by David Jones.
7
8'''
9  Example code integrating RGB PNG files, PyPNG and NumPy
10  (abstracted from Mel Raab's functioning code)
11'''
12
13# http://www.python.org/doc/2.4.4/lib/module-itertools.html
14import itertools
15
16import numpy
17import png
18
19
20''' If you have a PNG file for an RGB image,
21    and want to create a numpy array of data from it.
22'''
23# Read the file "picture.png" from the current directory.  The `Reader`
24# class can take a filename, a file-like object, or the byte data
25# directly; this suggests alternatives such as using urllib to read
26# an image from the internet:
27# png.Reader(file=urllib.urlopen('http://www.libpng.org/pub/png/PngSuite/basn2c16.png'))
28pngReader=png.Reader(filename='picture.png')
29# Tuple unpacking, using multiple assignment, is very useful for the
30# result of asDirect (and other methods).
31# See
32# http://docs.python.org/tutorial/introduction.html#first-steps-towards-programming
33row_count, column_count, pngdata, meta = pngReader.asDirect()
34bitdepth=meta['bitdepth']
35plane_count=meta['planes']
36
37# Make sure we're dealing with RGB files
38assert plane_count == 3
39
40''' Boxed row flat pixel:
41      list([R,G,B, R,G,B, R,G,B],
42           [R,G,B, R,G,B, R,G,B])
43    Array dimensions for this example:  (2,9)
44
45    Create `image_2d` as a two-dimensional NumPy array by stacking a
46    sequence of 1-dimensional arrays (rows).
47    The NumPy array mimics PyPNG's (boxed row flat pixel) representation;
48    it will have dimensions ``(row_count,column_count*plane_count)``.
49'''
50# The use of ``numpy.uint16``, below, is to convert each row to a NumPy
51# array with data type ``numpy.uint16``.  This is a feature of NumPy,
52# discussed further in
53# http://docs.scipy.org/doc/numpy/user/basics.types.html .
54# You can use avoid the explicit conversion with
55# ``numpy.vstack(pngdata)``, but then NumPy will pick the array's data
56# type; in practice it seems to pick ``numpy.int32``, which is large enough
57# to hold any pixel value for any PNG image but uses 4 bytes per value when
58# 1 or 2 would be enough.
59# --- extract 001 start
60image_2d = numpy.vstack(itertools.imap(numpy.uint16, pngdata))
61# --- extract 001 end
62# Do not be tempted to use ``numpy.asarray``; when passed an iterator
63# (`pngdata` is often an iterator) it will attempt to create a size 1
64# array with the iterator as its only element.
65# An alternative to the above is to create the target array of the right
66# shape, then populate it row by row:
67if 0:
68    image_2d = numpy.zeros((row_count,plane_count*column_count),
69                           dtype=numpy.uint16)
70    for row_index, one_boxed_row_flat_pixels in enumerate(pngdata):
71        image_2d[row_index,:]=one_boxed_row_flat_pixels
72
73del pngReader
74del pngdata
75
76
77''' Reconfigure for easier referencing, similar to
78        Boxed row boxed pixel:
79            list([ (R,G,B), (R,G,B), (R,G,B) ],
80                 [ (R,G,B), (R,G,B), (R,G,B) ])
81    Array dimensions for this example:  (2,3,3)
82
83    ``image_3d`` will contain the image as a three-dimensional numpy
84    array, having dimensions ``(row_count,column_count,plane_count)``.
85'''
86# --- extract 002 start
87image_3d = numpy.reshape(image_2d,
88                         (row_count,column_count,plane_count))
89# --- extract 002 end
90
91
92''' ============= '''
93
94''' Convert NumPy image_3d array to PNG image file.
95
96    If the data is three-dimensional, as it is above, the best thing
97    to do is reshape it into a two-dimensional array with a shape of
98    ``(row_count, column_count*plane_count)``.  Because a
99    two-dimensional numpy array is an iterator, it can be passed
100    directly to the ``png.Writer.write`` method.
101'''
102
103row_count, column_count, plane_count = image_3d.shape
104assert plane_count==3
105
106pngfile = open('picture_out.png', 'wb')
107try:
108    # This example assumes that you have 16-bit pixel values in the data
109    # array (that's what the ``bitdepth=16`` argument is for).
110    # If you don't, then the resulting PNG file will likely be
111    # very dark.  Hey, it's only an example.
112    pngWriter = png.Writer(column_count, row_count,
113                           greyscale=False,
114                           alpha=False,
115                           bitdepth=16)
116    # As of 2009-04-13 passing a numpy array that has an element type
117    # that is a numpy integer type (for example, the `image_3d` array has an
118    # element type of ``numpy.uint16``) generates a deprecation warning.
119    # This is probably a bug in numpy; it may go away in the future.
120    # The code still works despite the warning.
121    # See http://code.google.com/p/pypng/issues/detail?id=44
122# --- extract 003 start
123    pngWriter.write(pngfile,
124                    numpy.reshape(image_3d, (-1, column_count*plane_count)))
125# --- extract 003 end
126finally:
127    pngfile.close()
128
129