1#!/usr/bin/env python
2# -*- coding: utf-8 -*-
3# netpbmfile.py
4
5# Copyright (c) 2011-2013, Christoph Gohlke
6# Copyright (c) 2011-2013, The Regents of the University of California
7# Produced at the Laboratory for Fluorescence Dynamics.
8# All rights reserved.
9#
10# Redistribution and use in source and binary forms, with or without
11# modification, are permitted provided that the following conditions are met:
12#
13# * Redistributions of source code must retain the above copyright
14#   notice, this list of conditions and the following disclaimer.
15# * Redistributions in binary form must reproduce the above copyright
16#   notice, this list of conditions and the following disclaimer in the
17#   documentation and/or other materials provided with the distribution.
18# * Neither the name of the copyright holders nor the names of any
19#   contributors may be used to endorse or promote products derived
20#   from this software without specific prior written permission.
21#
22# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
23# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
24# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
25# ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
26# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
27# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
28# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
29# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
30# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
31# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
32# POSSIBILITY OF SUCH DAMAGE.
33
34"""Read and write image data from respectively to Netpbm files.
35
36This implementation follows the Netpbm format specifications at
37http://netpbm.sourceforge.net/doc/. No gamma correction is performed.
38
39The following image formats are supported: PBM (bi-level), PGM (grayscale),
40PPM (color), PAM (arbitrary), XV thumbnail (RGB332, read-only).
41
42:Author:
43  `Christoph Gohlke <http://www.lfd.uci.edu/~gohlke/>`_
44
45:Organization:
46  Laboratory for Fluorescence Dynamics, University of California, Irvine
47
48:Version: 2013.01.18
49
50Requirements
51------------
52* `CPython 2.7, 3.2 or 3.3 <http://www.python.org>`_
53* `Numpy 1.7 <http://www.numpy.org>`_
54* `Matplotlib 1.2 <http://www.matplotlib.org>`_  (optional for plotting)
55
56Examples
57--------
58>>> im1 = numpy.array([[0, 1],[65534, 65535]], dtype=numpy.uint16)
59>>> imsave('_tmp.pgm', im1)
60>>> im2 = imread('_tmp.pgm')
61>>> assert numpy.all(im1 == im2)
62
63"""
64
65from __future__ import division, print_function
66
67import sys
68import re
69import math
70from copy import deepcopy
71
72import numpy
73
74__version__ = '2013.01.18'
75__docformat__ = 'restructuredtext en'
76__all__ = ['imread', 'imsave', 'NetpbmFile']
77
78
79def imread(filename, *args, **kwargs):
80    """Return image data from Netpbm file as numpy array.
81
82    `args` and `kwargs` are arguments to NetpbmFile.asarray().
83
84    Examples
85    --------
86    >>> image = imread('_tmp.pgm')
87
88    """
89    try:
90        netpbm = NetpbmFile(filename)
91        image = netpbm.asarray()
92    finally:
93        netpbm.close()
94    return image
95
96
97def imsave(filename, data, maxval=None, pam=False):
98    """Write image data to Netpbm file.
99
100    Examples
101    --------
102    >>> image = numpy.array([[0, 1],[65534, 65535]], dtype=numpy.uint16)
103    >>> imsave('_tmp.pgm', image)
104
105    """
106    try:
107        netpbm = NetpbmFile(data, maxval=maxval)
108        netpbm.write(filename, pam=pam)
109    finally:
110        netpbm.close()
111
112
113class NetpbmFile(object):
114    """Read and write Netpbm PAM, PBM, PGM, PPM, files."""
115
116    _types = {b'P1': b'BLACKANDWHITE', b'P2': b'GRAYSCALE', b'P3': b'RGB',
117              b'P4': b'BLACKANDWHITE', b'P5': b'GRAYSCALE', b'P6': b'RGB',
118              b'P7 332': b'RGB', b'P7': b'RGB_ALPHA'}
119
120    def __init__(self, arg=None, **kwargs):
121        """Initialize instance from filename, open file, or numpy array."""
122        for attr in ('header', 'magicnum', 'width', 'height', 'maxval',
123                     'depth', 'tupltypes', '_filename', '_fh', '_data'):
124            setattr(self, attr, None)
125        if arg is None:
126            self._fromdata([], **kwargs)
127        elif isinstance(arg, basestring):
128            self._fh = open(arg, 'rb')
129            self._filename = arg
130            self._fromfile(self._fh, **kwargs)
131        elif hasattr(arg, 'seek'):
132            self._fromfile(arg, **kwargs)
133            self._fh = arg
134        else:
135            self._fromdata(arg, **kwargs)
136
137    def asarray(self, copy=True, cache=False, **kwargs):
138        """Return image data from file as numpy array."""
139        data = self._data
140        if data is None:
141            data = self._read_data(self._fh, **kwargs)
142            if cache:
143                self._data = data
144            else:
145                return data
146        return deepcopy(data) if copy else data
147
148    def write(self, arg, **kwargs):
149        """Write instance to file."""
150        if hasattr(arg, 'seek'):
151            self._tofile(arg, **kwargs)
152        else:
153            with open(arg, 'wb') as fid:
154                self._tofile(fid, **kwargs)
155
156    def close(self):
157        """Close open file. Future asarray calls might fail."""
158        if self._filename and self._fh:
159            self._fh.close()
160            self._fh = None
161
162    def __del__(self):
163        self.close()
164
165    def _fromfile(self, fh):
166        """Initialize instance from open file."""
167        fh.seek(0)
168        data = fh.read(4096)
169        if (len(data) < 7) or not (b'0' < data[1:2] < b'8'):
170            raise ValueError("Not a Netpbm file:\n%s" % data[:32])
171        try:
172            self._read_pam_header(data)
173        except Exception:
174            try:
175                self._read_pnm_header(data)
176            except Exception:
177                raise ValueError("Not a Netpbm file:\n%s" % data[:32])
178
179    def _read_pam_header(self, data):
180        """Read PAM header and initialize instance."""
181        regroups = re.search(
182            b"(^P7[\n\r]+(?:(?:[\n\r]+)|(?:#.*)|"
183            b"(HEIGHT\s+\d+)|(WIDTH\s+\d+)|(DEPTH\s+\d+)|(MAXVAL\s+\d+)|"
184            b"(?:TUPLTYPE\s+\w+))*ENDHDR\n)", data).groups()
185        self.header = regroups[0]
186        self.magicnum = b'P7'
187        for group in regroups[1:]:
188            key, value = group.split()
189            setattr(self, unicode(key).lower(), int(value))
190        matches = re.findall(b"(TUPLTYPE\s+\w+)", self.header)
191        self.tupltypes = [s.split(None, 1)[1] for s in matches]
192
193    def _read_pnm_header(self, data):
194        """Read PNM header and initialize instance."""
195        bpm = data[1:2] in b"14"
196        regroups = re.search(b"".join((
197            b"(^(P[123456]|P7 332)\s+(?:#.*[\r\n])*",
198            b"\s*(\d+)\s+(?:#.*[\r\n])*",
199            b"\s*(\d+)\s+(?:#.*[\r\n])*" * (not bpm),
200            b"\s*(\d+)\s(?:\s*#.*[\r\n]\s)*)")), data).groups() + (1, ) * bpm
201        self.header = regroups[0]
202        self.magicnum = regroups[1]
203        self.width = int(regroups[2])
204        self.height = int(regroups[3])
205        self.maxval = int(regroups[4])
206        self.depth = 3 if self.magicnum in b"P3P6P7 332" else 1
207        self.tupltypes = [self._types[self.magicnum]]
208
209    def _read_data(self, fh, byteorder='>'):
210        """Return image data from open file as numpy array."""
211        fh.seek(len(self.header))
212        data = fh.read()
213        dtype = 'u1' if self.maxval < 256 else byteorder + 'u2'
214        depth = 1 if self.magicnum == b"P7 332" else self.depth
215        shape = [-1, self.height, self.width, depth]
216        size = numpy.prod(shape[1:])
217        if self.magicnum in b"P1P2P3":
218            data = numpy.array(data.split(None, size)[:size], dtype)
219            data = data.reshape(shape)
220        elif self.maxval == 1:
221            shape[2] = int(math.ceil(self.width / 8))
222            data = numpy.frombuffer(data, dtype).reshape(shape)
223            data = numpy.unpackbits(data, axis=-2)[:, :, :self.width, :]
224        else:
225            data = numpy.frombuffer(data, dtype)
226            data = data[:size * (data.size // size)].reshape(shape)
227        if data.shape[0] < 2:
228            data = data.reshape(data.shape[1:])
229        if data.shape[-1] < 2:
230            data = data.reshape(data.shape[:-1])
231        if self.magicnum == b"P7 332":
232            rgb332 = numpy.array(list(numpy.ndindex(8, 8, 4)), numpy.uint8)
233            rgb332 *= [36, 36, 85]
234            data = numpy.take(rgb332, data, axis=0)
235        return data
236
237    def _fromdata(self, data, maxval=None):
238        """Initialize instance from numpy array."""
239        data = numpy.array(data, ndmin=2, copy=True)
240        if data.dtype.kind not in "uib":
241            raise ValueError("not an integer type: %s" % data.dtype)
242        if data.dtype.kind == 'i' and numpy.min(data) < 0:
243            raise ValueError("data out of range: %i" % numpy.min(data))
244        if maxval is None:
245            maxval = numpy.max(data)
246            maxval = 255 if maxval < 256 else 65535
247        if maxval < 0 or maxval > 65535:
248            raise ValueError("data out of range: %i" % maxval)
249        data = data.astype('u1' if maxval < 256 else '>u2')
250        self._data = data
251        if data.ndim > 2 and data.shape[-1] in (3, 4):
252            self.depth = data.shape[-1]
253            self.width = data.shape[-2]
254            self.height = data.shape[-3]
255            self.magicnum = b'P7' if self.depth == 4 else b'P6'
256        else:
257            self.depth = 1
258            self.width = data.shape[-1]
259            self.height = data.shape[-2]
260            self.magicnum = b'P5' if maxval > 1 else b'P4'
261        self.maxval = maxval
262        self.tupltypes = [self._types[self.magicnum]]
263        self.header = self._header()
264
265    def _tofile(self, fh, pam=False):
266        """Write Netbm file."""
267        fh.seek(0)
268        fh.write(self._header(pam))
269        data = self.asarray(copy=False)
270        if self.maxval == 1:
271            data = numpy.packbits(data, axis=-1)
272        data.tofile(fh)
273
274    def _header(self, pam=False):
275        """Return file header as byte string."""
276        if pam or self.magicnum == b'P7':
277            header = "\n".join((
278                "P7",
279                "HEIGHT %i" % self.height,
280                "WIDTH %i" % self.width,
281                "DEPTH %i" % self.depth,
282                "MAXVAL %i" % self.maxval,
283                "\n".join("TUPLTYPE %s" % unicode(i) for i in self.tupltypes),
284                "ENDHDR\n"))
285        elif self.maxval == 1:
286            header = "P4 %i %i\n" % (self.width, self.height)
287        elif self.depth == 1:
288            header = "P5 %i %i %i\n" % (self.width, self.height, self.maxval)
289        else:
290            header = "P6 %i %i %i\n" % (self.width, self.height, self.maxval)
291        if sys.version_info[0] > 2:
292            header = bytes(header, 'ascii')
293        return header
294
295    def __str__(self):
296        """Return information about instance."""
297        return unicode(self.header)
298
299
300if sys.version_info[0] > 2:
301    basestring = str
302    unicode = lambda x: str(x, 'ascii')
303
304if __name__ == "__main__":
305    # Show images specified on command line or all images in current directory
306    from glob import glob
307    from matplotlib import pyplot
308    files = sys.argv[1:] if len(sys.argv) > 1 else glob('*.p*m')
309    for fname in files:
310        try:
311            pam = NetpbmFile(fname)
312            img = pam.asarray(copy=False)
313            if False:
314                pam.write('_tmp.pgm.out', pam=True)
315                img2 = imread('_tmp.pgm.out')
316                assert numpy.all(img == img2)
317                imsave('_tmp.pgm.out', img)
318                img2 = imread('_tmp.pgm.out')
319                assert numpy.all(img == img2)
320            pam.close()
321        except ValueError as e:
322            print(fname, e)
323            continue
324        _shape = img.shape
325        if img.ndim > 3 or (img.ndim > 2 and img.shape[-1] not in (3, 4)):
326            img = img[0]
327        cmap = 'gray' if pam.maxval > 1 else 'binary'
328        pyplot.imshow(img, cmap, interpolation='nearest')
329        pyplot.title("%s %s %s %s" % (fname, unicode(pam.magicnum),
330                                      _shape, img.dtype))
331        pyplot.show()
332