home *** CD-ROM | disk | FTP | other *** search
/ Mac Easy 2010 May / Mac Life Ubuntu.iso / casper / filesystem.squashfs / usr / share / python-support / gnome-games-data / glchess / scene / opengl / png.py < prev    next >
Encoding:
Python Source  |  2009-04-14  |  38.6 KB  |  1,072 lines

  1. # -*- coding: utf-8 -*-
  2. #!/usr/bin/env python
  3. # png.py - PNG encoder in pure Python
  4. # Copyright (C) 2006 Johann C. Rocholl <johann@browsershots.org>
  5. #
  6. # Permission is hereby granted, free of charge, to any person
  7. # obtaining a copy of this software and associated documentation files
  8. # (the "Software"), to deal in the Software without restriction,
  9. # including without limitation the rights to use, copy, modify, merge,
  10. # publish, distribute, sublicense, and/or sell copies of the Software,
  11. # and to permit persons to whom the Software is furnished to do so,
  12. # subject to the following conditions:
  13. #
  14. # The above copyright notice and this permission notice shall be
  15. # included in all copies or substantial portions of the Software.
  16. #
  17. # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
  18. # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
  19. # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
  20. # NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
  21. # BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
  22. # ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
  23. # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
  24. # SOFTWARE.
  25. #
  26. # Contributors (alphabetical):
  27. # Nicko van Someren <nicko@nicko.org>
  28. #
  29. # Changelog (recent first):
  30. # 2006-06-17 Nicko: Reworked into a class, faster interlacing.
  31. # 2006-06-17 Johann: Very simple prototype PNG decoder.
  32. # 2006-06-17 Nicko: Test suite with various image generators.
  33. # 2006-06-17 Nicko: Alpha-channel, grey-scale, 16-bit/plane support.
  34. # 2006-06-15 Johann: Scanline iterator interface for large input files.
  35. # 2006-06-09 Johann: Very simple prototype PNG encoder.
  36.  
  37.  
  38. """
  39. Pure Python PNG Reader/Writer
  40.  
  41. This is an implementation of a subset of the PNG specification at
  42. http://www.w3.org/TR/2003/REC-PNG-20031110 in pure Python. It reads
  43. and writes PNG files with 8/16/24/32/48/64 bits per pixel (greyscale,
  44. RGB, RGBA, with 8 or 16 bits per layer), with a number of options. For
  45. help, type "import png; help(png)" in your python interpreter.
  46.  
  47. This file can also be used as a command-line utility to convert PNM
  48. files to PNG. The interface is similar to that of the pnmtopng program
  49. from the netpbm package. Type "python png.py --help" at the shell
  50. prompt for usage and a list of options.
  51. """
  52.  
  53.  
  54. __revision__ = '$Rev$'
  55. __date__ = '$Date$'
  56. __author__ = '$Author$'
  57.  
  58.  
  59. import sys
  60. import zlib
  61. import struct
  62. import math
  63. from array import array
  64.  
  65.  
  66. _adam7 = ((0, 0, 8, 8),
  67.           (4, 0, 8, 8),
  68.           (0, 4, 4, 8),
  69.           (2, 0, 4, 4),
  70.           (0, 2, 2, 4),
  71.           (1, 0, 2, 2),
  72.           (0, 1, 1, 2))
  73.  
  74.  
  75. def interleave_planes(ipixels, apixels, ipsize, apsize):
  76.     """
  77.     Interleave color planes, e.g. RGB + A = RGBA.
  78.  
  79.     Return an array of pixels consisting of the ipsize bytes of data
  80.     from each pixel in ipixels followed by the apsize bytes of data
  81.     from each pixel in apixels, for an image of size width x height.
  82.     """
  83.     itotal = len(ipixels)
  84.     atotal = len(apixels)
  85.     newtotal = itotal + atotal
  86.     newpsize = ipsize + apsize
  87.     # Set up the output buffer
  88.     out = array('B')
  89.     # It's annoying that there is no cheap way to set the array size :-(
  90.     out.extend(ipixels)
  91.     out.extend(apixels)
  92.     # Interleave in the pixel data
  93.     for i in range(ipsize):
  94.         out[i:newtotal:newpsize] = ipixels[i:itotal:ipsize]
  95.     for i in range(apsize):
  96.         out[i+ipsize:newtotal:newpsize] = apixels[i:atotal:apsize]
  97.     return out
  98.  
  99. class Error(Exception):
  100.     pass
  101.  
  102. class Writer:
  103.     """
  104.     PNG encoder in pure Python.
  105.     """
  106.  
  107.     def __init__(self, width, height,
  108.                  transparent=None,
  109.                  background=None,
  110.                  gamma=None,
  111.                  greyscale=False,
  112.                  has_alpha=False,
  113.                  bytes_per_sample=1,
  114.                  compression=None,
  115.                  interlaced=False,
  116.                  chunk_limit=2**20):
  117.         """
  118.         Create a PNG encoder object.
  119.  
  120.         Arguments:
  121.         width, height - size of the image in pixels
  122.         transparent - create a tRNS chunk
  123.         background - create a bKGD chunk
  124.         gamma - create a gAMA chunk
  125.         greyscale - input data is greyscale, not RGB
  126.         has_alpha - input data has alpha channel (RGBA)
  127.         bytes_per_sample - 8-bit or 16-bit input data
  128.         compression - zlib compression level (1-9)
  129.         chunk_limit - write multiple IDAT chunks to save memory
  130.  
  131.         If specified, the transparent and background parameters must
  132.         be a tuple with three integer values for red, green, blue, or
  133.         a simple integer (or singleton tuple) for a greyscale image.
  134.  
  135.         If specified, the gamma parameter must be a float value.
  136.  
  137.         """
  138.         if width <= 0 or height <= 0:
  139.             raise ValueError("width and height must be greater than zero")
  140.  
  141.         if has_alpha and transparent is not None:
  142.             raise ValueError(
  143.                 "transparent color not allowed with alpha channel")
  144.  
  145.         if bytes_per_sample < 1 or bytes_per_sample > 2:
  146.             raise ValueError("bytes per sample must be 1 or 2")
  147.  
  148.         if transparent is not None:
  149.             if greyscale:
  150.                 if type(transparent) is not int:
  151.                     raise ValueError(
  152.                         "transparent color for greyscale must be integer")
  153.             else:
  154.                 if not (len(transparent) == 3 and
  155.                         type(transparent[0]) is int and
  156.                         type(transparent[1]) is int and
  157.                         type(transparent[2]) is int):
  158.                     raise ValueError(
  159.                         "transparent color must be a triple of integers")
  160.  
  161.         if background is not None:
  162.             if greyscale:
  163.                 if type(background) is not int:
  164.                     raise ValueError(
  165.                         "background color for greyscale must be integer")
  166.             else:
  167.                 if not (len(background) == 3 and
  168.                         type(background[0]) is int and
  169.                         type(background[1]) is int and
  170.                         type(background[2]) is int):
  171.                     raise ValueError(
  172.                         "background color must be a triple of integers")
  173.  
  174.         self.width = width
  175.         self.height = height
  176.         self.transparent = transparent
  177.         self.background = background
  178.         self.gamma = gamma
  179.         self.greyscale = greyscale
  180.         self.has_alpha = has_alpha
  181.         self.bytes_per_sample = bytes_per_sample
  182.         self.compression = compression
  183.         self.chunk_limit = chunk_limit
  184.         self.interlaced = interlaced
  185.  
  186.         if self.greyscale:
  187.             self.color_depth = 1
  188.             if self.has_alpha:
  189.                 self.color_type = 4
  190.                 self.psize = self.bytes_per_sample * 2
  191.             else:
  192.                 self.color_type = 0
  193.                 self.psize = self.bytes_per_sample
  194.         else:
  195.             self.color_depth = 3
  196.             if self.has_alpha:
  197.                 self.color_type = 6
  198.                 self.psize = self.bytes_per_sample * 4
  199.             else:
  200.                 self.color_type = 2
  201.                 self.psize = self.bytes_per_sample * 3
  202.  
  203.     def write_chunk(self, outfile, tag, data):
  204.         """
  205.         Write a PNG chunk to the output file, including length and checksum.
  206.         """
  207.         # http://www.w3.org/TR/PNG/#5Chunk-layout
  208.         outfile.write(struct.pack("!I", len(data)))
  209.         outfile.write(tag)
  210.         outfile.write(data)
  211.         checksum = zlib.crc32(tag)
  212.         checksum = zlib.crc32(data, checksum)
  213.         outfile.write(struct.pack("!i", checksum))
  214.  
  215.     def write(self, outfile, scanlines):
  216.         """
  217.         Write a PNG image to the output file.
  218.         """
  219.         # http://www.w3.org/TR/PNG/#5PNG-file-signature
  220.         outfile.write(struct.pack("8B", 137, 80, 78, 71, 13, 10, 26, 10))
  221.  
  222.         # http://www.w3.org/TR/PNG/#11IHDR
  223.         if self.interlaced:
  224.             interlaced = 1
  225.         else:
  226.             interlaced = 0
  227.         self.write_chunk(outfile, 'IHDR',
  228.                          struct.pack("!2I5B", self.width, self.height,
  229.                                      self.bytes_per_sample * 8,
  230.                                      self.color_type, 0, 0, interlaced))
  231.  
  232.         # http://www.w3.org/TR/PNG/#11tRNS
  233.         if self.transparent is not None:
  234.             if self.greyscale:
  235.                 self.write_chunk(outfile, 'tRNS',
  236.                                  struct.pack("!1H", *self.transparent))
  237.             else:
  238.                 self.write_chunk(outfile, 'tRNS',
  239.                                  struct.pack("!3H", *self.transparent))
  240.  
  241.         # http://www.w3.org/TR/PNG/#11bKGD
  242.         if self.background is not None:
  243.             if self.greyscale:
  244.                 self.write_chunk(outfile, 'bKGD',
  245.                                  struct.pack("!1H", *self.background))
  246.             else:
  247.                 self.write_chunk(outfile, 'bKGD',
  248.                                  struct.pack("!3H", *self.background))
  249.  
  250.         # http://www.w3.org/TR/PNG/#11gAMA
  251.         if self.gamma is not None:
  252.             self.write_chunk(outfile, 'gAMA',
  253.                              struct.pack("!L", int(self.gamma * 100000)))
  254.  
  255.         # http://www.w3.org/TR/PNG/#11IDAT
  256.         if self.compression is not None:
  257.             compressor = zlib.compressobj(self.compression)
  258.         else:
  259.             compressor = zlib.compressobj()
  260.  
  261.         data = array('B')
  262.         for scanline in scanlines:
  263.             data.append(0)
  264.             data.extend(scanline)
  265.             if len(data) > self.chunk_limit:
  266.                 compressed = compressor.compress(data.tostring())
  267.                 if len(compressed):
  268.                     # print >> sys.stderr, len(data), len(compressed)
  269.                     self.write_chunk(outfile, 'IDAT', compressed)
  270.                 data = array('B')
  271.         if len(data):
  272.             compressed = compressor.compress(data.tostring())
  273.         else:
  274.             compressed = ''
  275.         flushed = compressor.flush()
  276.         if len(compressed) or len(flushed):
  277.             # print >> sys.stderr, len(data), len(compressed), len(flushed)
  278.             self.write_chunk(outfile, 'IDAT', compressed + flushed)
  279.  
  280.         # http://www.w3.org/TR/PNG/#11IEND
  281.         self.write_chunk(outfile, 'IEND', '')
  282.  
  283.     def write_array(self, outfile, pixels):
  284.         """
  285.         Encode a pixel array to PNG and write output file.
  286.         """
  287.         if self.interlaced:
  288.             self.write(outfile, self.array_scanlines_interlace(pixels))
  289.         else:
  290.             self.write(outfile, self.array_scanlines(pixels))
  291.  
  292.     def convert_ppm(self, ppmfile, outfile):
  293.         """
  294.         Convert a PPM file containing raw pixel data into a PNG file
  295.         with the parameters set in the writer object.
  296.         """
  297.         if self.interlaced:
  298.             pixels = array('B')
  299.             pixels.fromfile(ppmfile,
  300.                             self.bytes_per_sample * self.color_depth *
  301.                             self.width * self.height)
  302.             self.write(outfile, self.array_scanlines_interlace(pixels))
  303.         else:
  304.             self.write(outfile, self.file_scanlines(ppmfile))
  305.  
  306.     def convert_ppm_and_pgm(self, ppmfile, pgmfile, outfile):
  307.         """
  308.         Convert a PPM and PGM file containing raw pixel data into a
  309.         PNG outfile with the parameters set in the writer object.
  310.         """
  311.         pixels = array('B')
  312.         pixels.fromfile(ppmfile,
  313.                         self.bytes_per_sample * self.color_depth *
  314.                         self.width * self.height)
  315.         apixels = array('B')
  316.         apixels.fromfile(pgmfile,
  317.                          self.bytes_per_sample *
  318.                          self.width * self.height)
  319.         pixels = interleave_planes(pixels, apixels,
  320.                                    self.bytes_per_sample * self.color_depth,
  321.                                    self.bytes_per_sample)
  322.         if self.interlaced:
  323.             self.write(outfile, self.array_scanlines_interlace(pixels))
  324.         else:
  325.             self.write(outfile, self.array_scanlines(pixels))
  326.  
  327.     def file_scanlines(self, infile):
  328.         """
  329.         Generator for scanlines from an input file.
  330.         """
  331.         row_bytes = self.psize * self.width
  332.         for y in range(self.height):
  333.             scanline = array('B')
  334.             scanline.fromfile(infile, row_bytes)
  335.             yield scanline
  336.  
  337.     def array_scanlines(self, pixels):
  338.         """
  339.         Generator for scanlines from an array.
  340.         """
  341.         row_bytes = self.width * self.psize
  342.         stop = 0
  343.         for y in range(self.height):
  344.             start = stop
  345.             stop = start + row_bytes
  346.             yield pixels[start:stop]
  347.  
  348.     def old_array_scanlines_interlace(self, pixels):
  349.         """
  350.         Generator for interlaced scanlines from an array.
  351.         http://www.w3.org/TR/PNG/#8InterlaceMethods
  352.         """
  353.         row_bytes = self.psize * self.width
  354.         for xstart, ystart, xstep, ystep in _adam7:
  355.             for y in range(ystart, self.height, ystep):
  356.                 if xstart < self.width:
  357.                     if xstep == 1:
  358.                         offset = y*row_bytes
  359.                         yield pixels[offset:offset+row_bytes]
  360.                     else:
  361.                         row = array('B')
  362.                         offset = y*row_bytes + xstart* self.psize
  363.                         skip = self.psize * xstep
  364.                         for x in range(xstart, self.width, xstep):
  365.                             row.extend(pixels[offset:offset + self.psize])
  366.                             offset += skip
  367.                         yield row
  368.  
  369.     def array_scanlines_interlace(self, pixels):
  370.         """
  371.         Generator for interlaced scanlines from an array.
  372.         http://www.w3.org/TR/PNG/#8InterlaceMethods
  373.         """
  374.         row_bytes = self.psize * self.width
  375.         for xstart, ystart, xstep, ystep in _adam7:
  376.             for y in range(ystart, self.height, ystep):
  377.                 if xstart >= self.width:
  378.                     continue
  379.                 if xstep == 1:
  380.                     offset = y * row_bytes
  381.                     yield pixels[offset:offset+row_bytes]
  382.                 else:
  383.                     row = array('B')
  384.                     # Note we want the ceiling of (self.width - xstart) / xtep
  385.                     row_len = self.psize * (
  386.                         (self.width - xstart + xstep - 1) / xstep)
  387.                     # There's no easier way to set the length of an array
  388.                     row.extend(pixels[0:row_len])
  389.                     offset = y * row_bytes + xstart * self.psize
  390.                     end_offset = (y+1) * row_bytes
  391.                     skip = self.psize * xstep
  392.                     for i in range(self.psize):
  393.                         row[i:row_len:self.psize] = \
  394.                             pixels[offset+i:end_offset:skip]
  395.                     yield row
  396.  
  397. class _readable:
  398.     """
  399.     A simple file-like interface for strings and arrays.
  400.     """
  401.  
  402.     def __init__(self, buf):
  403.         self.buf = buf
  404.         self.offset = 0
  405.  
  406.     def read(self, n):
  407.         r = self.buf[self.offset:self.offset+n]
  408.         if isinstance(r, array):
  409.             r = r.tostring()
  410.         self.offset += n
  411.         return r
  412.  
  413. class Reader:
  414.     """
  415.     PNG decoder in pure Python.
  416.     """
  417.  
  418.     def __init__(self, _guess=None, **kw):
  419.         """
  420.         Create a PNG decoder object.
  421.  
  422.         The constructor expects exactly one keyword argument. If you
  423.         supply a positional argument instead, it will guess the input
  424.         type. You can choose among the following arguments:
  425.         filename - name of PNG input file
  426.         file - object with a read() method
  427.         pixels - array or string with PNG data
  428.  
  429.         """
  430.         if ((_guess is not None and len(kw) != 0) or
  431.             (_guess is None and len(kw) != 1)):
  432.             raise TypeError("Reader() takes exactly 1 argument")
  433.  
  434.         if _guess is not None:
  435.             if isinstance(_guess, array):
  436.                 kw["pixels"] = _guess
  437.             elif isinstance(_guess, str):
  438.                 kw["filename"] = _guess
  439.             elif isinstance(_guess, file):
  440.                 kw["file"] = _guess
  441.  
  442.         if "filename" in kw:
  443.             self.file = file(kw["filename"], "rb")
  444.         elif "file" in kw:
  445.             self.file = kw["file"]
  446.         elif "pixels" in kw:
  447.             self.file = _readable(kw["pixels"])
  448.         else:
  449.             raise TypeError("expecting filename, file or pixels array")
  450.  
  451.     def read_chunk(self):
  452.         """
  453.         Read a PNG chunk from the input file, return tag name and data.
  454.         """
  455.         # http://www.w3.org/TR/PNG/#5Chunk-layout
  456.         try:
  457.             data_bytes, tag = struct.unpack('!I4s', self.file.read(8))
  458.         except struct.error:
  459.             raise ValueError('Chunk too short for header')
  460.         data = self.file.read(data_bytes)
  461.         if len(data) != data_bytes:
  462.             raise ValueError('Chunk %s too short for required %i data octets'
  463.                              % (tag, data_bytes))
  464.         checksum = self.file.read(4)
  465.         if len(checksum) != 4:
  466.             raise ValueError('Chunk %s too short for checksum', tag)
  467.         verify = zlib.crc32(tag)
  468.         verify = zlib.crc32(data, verify)
  469.         verify = struct.pack('!i', verify)
  470.         if checksum != verify:
  471.             # print repr(checksum)
  472.             (a, ) = struct.unpack('!I', checksum)
  473.             (b, ) = struct.unpack('!I', verify)
  474.             raise ValueError("Checksum error in %s chunk: 0x%X != 0x%X"
  475.                              % (tag, a, b))
  476.         return tag, data
  477.  
  478.     def _reconstruct_sub(self, offset, xstep, ystep):
  479.         """
  480.         Reverse sub filter.
  481.         """
  482.         pixels = self.pixels
  483.         a_offset = offset
  484.         offset += self.psize * xstep
  485.         if xstep == 1:
  486.             for index in range(self.psize, self.row_bytes):
  487.                 x = pixels[offset]
  488.                 a = pixels[a_offset]
  489.                 pixels[offset] = (x + a) & 0xff
  490.                 offset += 1
  491.                 a_offset += 1
  492.         else:
  493.             byte_step = self.psize * xstep
  494.             for index in range(byte_step, self.row_bytes, byte_step):
  495.                 for i in range(self.psize):
  496.                     x = pixels[offset + i]
  497.                     a = pixels[a_offset + i]
  498.                     pixels[offset + i] = (x + a) & 0xff
  499.                 offset += self.psize * xstep
  500.                 a_offset += self.psize * xstep
  501.  
  502.     def _reconstruct_up(self, offset, xstep, ystep):
  503.         """
  504.         Reverse up filter.
  505.         """
  506.         pixels = self.pixels
  507.         b_offset = offset - (self.row_bytes * ystep)
  508.         if xstep == 1:
  509.             for index in range(self.row_bytes):
  510.                 x = pixels[offset]
  511.                 b = pixels[b_offset]
  512.                 pixels[offset] = (x + b) & 0xff
  513.                 offset += 1
  514.                 b_offset += 1
  515.         else:
  516.             for index in range(0, self.row_bytes, xstep * self.psize):
  517.                 for i in range(self.psize):
  518.                     x = pixels[offset + i]
  519.                     b = pixels[b_offset + i]
  520.                     pixels[offset + i] = (x + b) & 0xff
  521.                 offset += self.psize * xstep
  522.                 b_offset += self.psize * xstep
  523.  
  524.     def _reconstruct_average(self, offset, xstep, ystep):
  525.         """
  526.         Reverse average filter.
  527.         """
  528.         pixels = self.pixels
  529.         a_offset = offset - (self.psize * xstep)
  530.         b_offset = offset - (self.row_bytes * ystep)
  531.         if xstep == 1:
  532.             for index in range(self.row_bytes):
  533.                 x = pixels[offset]
  534.                 if index < self.psize:
  535.                     a = 0
  536.                 else:
  537.                     a = pixels[a_offset]
  538.                 if b_offset < 0:
  539.                     b = 0
  540.                 else:
  541.                     b = pixels[b_offset]
  542.                 pixels[offset] = (x + ((a + b) >> 1)) & 0xff
  543.                 offset += 1
  544.                 a_offset += 1
  545.                 b_offset += 1
  546.         else:
  547.             for index in range(0, self.row_bytes, self.psize * xstep):
  548.                 for i in range(self.psize):
  549.                     x = pixels[offset+i]
  550.                     if index < self.psize:
  551.                         a = 0
  552.                     else:
  553.                         a = pixels[a_offset + i]
  554.                     if b_offset < 0:
  555.                         b = 0
  556.                     else:
  557.                         b = pixels[b_offset + i]
  558.                     pixels[offset + i] = (x + ((a + b) >> 1)) & 0xff
  559.                 offset += self.psize * xstep
  560.                 a_offset += self.psize * xstep
  561.                 b_offset += self.psize * xstep
  562.  
  563.     def _reconstruct_paeth(self, offset, xstep, ystep):
  564.         """
  565.         Reverse Paeth filter.
  566.         """
  567.         pixels = self.pixels
  568.         a_offset = offset - (self.psize * xstep)
  569.         b_offset = offset - (self.row_bytes * ystep)
  570.         c_offset = b_offset - (self.psize * xstep)
  571.         # There's enough inside this loop that it's probably not worth
  572.         # optimising for xstep == 1
  573.         for index in range(0, self.row_bytes, self.psize * xstep):
  574.             for i in range(self.psize):
  575.                 x = pixels[offset+i]
  576.                 if index < self.psize:
  577.                     a = c = 0
  578.                     b = pixels[b_offset+i]
  579.                 else:
  580.                     a = pixels[a_offset+i]
  581.                     b = pixels[b_offset+i]
  582.                     c = pixels[c_offset+i]
  583.                 p = a + b - c
  584.                 pa = abs(p - a)
  585.                 pb = abs(p - b)
  586.                 pc = abs(p - c)
  587.                 if pa <= pb and pa <= pc:
  588.                     pr = a
  589.                 elif pb <= pc:
  590.                     pr = b
  591.                 else:
  592.                     pr = c
  593.                 pixels[offset+i] = (x + pr) & 0xff
  594.             offset += self.psize * xstep
  595.             a_offset += self.psize * xstep
  596.             b_offset += self.psize * xstep
  597.             c_offset += self.psize * xstep
  598.  
  599.     # N.B. PNG files with 'up', 'average' or 'paeth' filters on the
  600.     # first line of a pass are legal. The code above for 'average'
  601.     # deals with this case explicitly. For up we map to the null
  602.     # filter and for paeth we map to the sub filter.
  603.  
  604.     def reconstruct_line(self, filter_type, first_line, offset, xstep, ystep):
  605.         """
  606.         Reverse the filtering for a scanline.
  607.         """
  608.         # print >> sys.stderr, "Filter type %s, first_line=%s" % (
  609.         #                      filter_type, first_line)
  610.         filter_type += (first_line << 8)
  611.         if filter_type == 1 or filter_type == 0x101 or filter_type == 0x104:
  612.             self._reconstruct_sub(offset, xstep, ystep)
  613.         elif filter_type == 2:
  614.             self._reconstruct_up(offset, xstep, ystep)
  615.         elif filter_type == 3 or filter_type == 0x103:
  616.             self._reconstruct_average(offset, xstep, ystep)
  617.         elif filter_type == 4:
  618.             self._reconstruct_paeth(offset, xstep, ystep)
  619.         return
  620.  
  621.     def deinterlace(self, scanlines):
  622.         """
  623.         Read pixel data and remove interlacing.
  624.         """
  625.         # print >> sys.stderr, ("Reading interlaced, w=%s, r=%s, planes=%s," +
  626.         #     " bpp=%s") % (self.width, self.height, self.planes, self.bps)
  627.         a = array('B')
  628.         self.pixels = a
  629.         # Make the array big enough
  630.         temp = scanlines[0:self.width*self.height*self.psize]
  631.         a.extend(temp)
  632.         source_offset = 0
  633.         for xstart, ystart, xstep, ystep in _adam7:
  634.             # print >> sys.stderr, "Adam7: start=%s,%s step=%s,%s" % (
  635.             #     xstart, ystart, xstep, ystep)
  636.             filter_first_line = 1
  637.             for y in range(ystart, self.height, ystep):
  638.                 if xstart >= self.width:
  639.                     continue
  640.                 filter_type = scanlines[source_offset]
  641.                 source_offset += 1
  642.                 if xstep == 1:
  643.                     offset = y * self.row_bytes
  644.                     a[offset:offset+self.row_bytes] = \
  645.                         scanlines[source_offset:source_offset + self.row_bytes]
  646.                     source_offset += self.row_bytes
  647.                 else:
  648.                     # Note we want the ceiling of (width - xstart) / xtep
  649.                     row_len = self.psize * (
  650.                         (self.width - xstart + xstep - 1) / xstep)
  651.                     offset = y * self.row_bytes + xstart * self.psize
  652.                     end_offset = (y+1) * self.row_bytes
  653.                     skip = self.psize * xstep
  654.                     for i in range(self.psize):
  655.                         a[offset+i:end_offset:skip] = \
  656.                             scanlines[source_offset + i:
  657.                                       source_offset + row_len:
  658.                                       self.psize]
  659.                     source_offset += row_len
  660.                 if filter_type:
  661.                     self.reconstruct_line(filter_type, filter_first_line,
  662.                                           offset, xstep, ystep)
  663.                 filter_first_line = 0
  664.         return a
  665.  
  666.     def read_flat(self, scanlines):
  667.         """
  668.         Read pixel data without de-interlacing.
  669.         """
  670.         a = array('B')
  671.         self.pixels = a
  672.         offset = 0
  673.         source_offset = 0
  674.         filter_first_line = 1
  675.         for y in range(self.height):
  676.             filter_type = scanlines[source_offset]
  677.             source_offset += 1
  678.             a.extend(scanlines[source_offset: source_offset + self.row_bytes])
  679.             if filter_type:
  680.                 self.reconstruct_line(filter_type, filter_first_line,
  681.                                       offset, 1, 1)
  682.             filter_first_line = 0
  683.             offset += self.row_bytes
  684.             source_offset += self.row_bytes
  685.         return a
  686.  
  687.     def read(self):
  688.         """
  689.         Read a simple PNG file, return width, height, pixels and image metadata
  690.  
  691.         This function is a very early prototype with limited flexibility
  692.         and excessive use of memory.
  693.         """
  694.         signature = self.file.read(8)
  695.         if (signature != struct.pack("8B", 137, 80, 78, 71, 13, 10, 26, 10)):
  696.             raise Error("PNG file has invalid header")
  697.         compressed = []
  698.         image_metadata = {}
  699.         while True:
  700.             try:
  701.                 tag, data = self.read_chunk()
  702.             except ValueError, e:
  703.                 raise Error('Chunk error: ' + e.message)
  704.  
  705.             # print >> sys.stderr, tag, len(data)
  706.             if tag == 'IHDR': # http://www.w3.org/TR/PNG/#11IHDR
  707.                 (width, height, bits_per_sample, color_type,
  708.                  compression_method, filter_method,
  709.                  interlaced) = struct.unpack("!2I5B", data)
  710.                 bps = bits_per_sample / 8
  711.                 if bps == 0:
  712.                     raise Error("unsupported pixel depth")
  713.                 if bps > 2 or bits_per_sample != (bps * 8):
  714.                     raise Error("invalid pixel depth")
  715.                 if color_type == 0:
  716.                     greyscale = True
  717.                     has_alpha = False
  718.                     planes = 1
  719.                 elif color_type == 2:
  720.                     greyscale = False
  721.                     has_alpha = False
  722.                     planes = 3
  723.                 elif color_type == 4:
  724.                     greyscale = True
  725.                     has_alpha = True
  726.                     planes = 2
  727.                 elif color_type == 6:
  728.                     greyscale = False
  729.                     has_alpha = True
  730.                     planes = 4
  731.                 else:
  732.                     raise Error("unknown PNG colour type %s" % color_type)
  733.                 if compression_method != 0:
  734.                     raise Error("unknown compression method")
  735.                 if filter_method != 0:
  736.                     raise Error("unknown filter method")
  737.                 self.bps = bps
  738.                 self.planes = planes
  739.                 self.psize = bps * planes
  740.                 self.width = width
  741.                 self.height = height
  742.                 self.row_bytes = width * self.psize
  743.             elif tag == 'IDAT': # http://www.w3.org/TR/PNG/#11IDAT
  744.                 compressed.append(data)
  745.             elif tag == 'bKGD':
  746.                 if greyscale:
  747.                     image_metadata["background"] = struct.unpack("!1H", data)
  748.                 else:
  749.                     image_metadata["background"] = struct.unpack("!3H", data)
  750.             elif tag == 'tRNS':
  751.                 if greyscale:
  752.                     image_metadata["transparent"] = struct.unpack("!1H", data)
  753.                 else:
  754.                     image_metadata["transparent"] = struct.unpack("!3H", data)
  755.             elif tag == 'gAMA':
  756.                 image_metadata["gamma"] = (
  757.                     struct.unpack("!L", data)[0]) / 100000.0
  758.             elif tag == 'IEND': # http://www.w3.org/TR/PNG/#11IEND
  759.                 break
  760.         scanlines = array('B', zlib.decompress(''.join(compressed)))
  761.         if interlaced:
  762.             pixels = self.deinterlace(scanlines)
  763.         else:
  764.             pixels = self.read_flat(scanlines)
  765.         image_metadata["greyscale"] = greyscale
  766.         image_metadata["has_alpha"] = has_alpha
  767.         image_metadata["bytes_per_sample"] = bps
  768.         image_metadata["interlaced"] = interlaced
  769.         return width, height, pixels, image_metadata
  770.  
  771.  
  772. def test_suite(options):
  773.     """
  774.     Run regression test and write PNG file to stdout.
  775.     """
  776.  
  777.     # Below is a big stack of test image generators
  778.  
  779.     def test_gradient_horizontal_lr(x, y):
  780.         return x
  781.  
  782.     def test_gradient_horizontal_rl(x, y):
  783.         return 1-x
  784.  
  785.     def test_gradient_vertical_tb(x, y):
  786.         return y
  787.  
  788.     def test_gradient_vertical_bt(x, y):
  789.         return 1-y
  790.  
  791.     def test_radial_tl(x, y):
  792.         return max(1-math.sqrt(x*x+y*y), 0.0)
  793.  
  794.     def test_radial_center(x, y):
  795.         return test_radial_tl(x-0.5, y-0.5)
  796.  
  797.     def test_radial_tr(x, y):
  798.         return test_radial_tl(1-x, y)
  799.  
  800.     def test_radial_bl(x, y):
  801.         return test_radial_tl(x, 1-y)
  802.  
  803.     def test_radial_br(x, y):
  804.         return test_radial_tl(1-x, 1-y)
  805.  
  806.     def test_stripe(x, n):
  807.         return 1.0*(int(x*n) & 1)
  808.  
  809.     def test_stripe_h_2(x, y):
  810.         return test_stripe(x, 2)
  811.  
  812.     def test_stripe_h_4(x, y):
  813.         return test_stripe(x, 4)
  814.  
  815.     def test_stripe_h_10(x, y):
  816.         return test_stripe(x, 10)
  817.  
  818.     def test_stripe_v_2(x, y):
  819.         return test_stripe(y, 2)
  820.  
  821.     def test_stripe_v_4(x, y):
  822.         return test_stripe(y, 4)
  823.  
  824.     def test_stripe_v_10(x, y):
  825.         return test_stripe(y, 10)
  826.  
  827.     def test_stripe_lr_10(x, y):
  828.         return test_stripe(x+y, 10)
  829.  
  830.     def test_stripe_rl_10(x, y):
  831.         return test_stripe(x-y, 10)
  832.  
  833.     def test_checker(x, y, n):
  834.         return 1.0*((int(x*n) & 1) ^ (int(y*n) & 1))
  835.  
  836.     def test_checker_8(x, y):
  837.         return test_checker(x, y, 8)
  838.  
  839.     def test_checker_15(x, y):
  840.         return test_checker(x, y, 15)
  841.  
  842.     def test_zero(x, y):
  843.         return 0
  844.  
  845.     def test_one(x, y):
  846.         return 1
  847.  
  848.     test_patterns = {
  849.         "GLR": test_gradient_horizontal_lr,
  850.         "GRL": test_gradient_horizontal_rl,
  851.         "GTB": test_gradient_vertical_tb,
  852.         "GBT": test_gradient_vertical_bt,
  853.         "RTL": test_radial_tl,
  854.         "RTR": test_radial_tr,
  855.         "RBL": test_radial_bl,
  856.         "RBR": test_radial_br,
  857.         "RCTR": test_radial_center,
  858.         "HS2": test_stripe_h_2,
  859.         "HS4": test_stripe_h_4,
  860.         "HS10": test_stripe_h_10,
  861.         "VS2": test_stripe_v_2,
  862.         "VS4": test_stripe_v_4,
  863.         "VS10": test_stripe_v_10,
  864.         "LRS": test_stripe_lr_10,
  865.         "RLS": test_stripe_rl_10,
  866.         "CK8": test_checker_8,
  867.         "CK15": test_checker_15,
  868.         "ZERO": test_zero,
  869.         "ONE": test_one,
  870.         }
  871.  
  872.     def test_pattern(width, height, depth, pattern):
  873.         """
  874.         Create a single plane (monochrome) test pattern.
  875.         """
  876.         a = array('B')
  877.         fw = float(width)
  878.         fh = float(height)
  879.         pfun = test_patterns[pattern]
  880.         if depth == 1:
  881.             for y in range(height):
  882.                 for x in range(width):
  883.                     a.append(int(pfun(float(x)/fw, float(y)/fh) * 255))
  884.         elif depth == 2:
  885.             for y in range(height):
  886.                 for x in range(width):
  887.                     v = int(pfun(float(x)/fw, float(y)/fh) * 65535)
  888.                     a.append(v >> 8)
  889.                     a.append(v & 0xff)
  890.         return a
  891.  
  892.     def test_rgba(size=256, depth=1,
  893.                     red="GTB", green="GLR", blue="RTL", alpha=None):
  894.         """
  895.         Create a test image.
  896.         """
  897.         r = test_pattern(size, size, depth, red)
  898.         g = test_pattern(size, size, depth, green)
  899.         b = test_pattern(size, size, depth, blue)
  900.         if alpha:
  901.             a = test_pattern(size, size, depth, alpha)
  902.         i = interleave_planes(r, g, depth, depth)
  903.         i = interleave_planes(i, b, 2 * depth, depth)
  904.         if alpha:
  905.             i = interleave_planes(i, a, 3 * depth, depth)
  906.         return i
  907.  
  908.     # The body of test_suite()
  909.     size = 256
  910.     if options.test_size:
  911.         size = options.test_size
  912.     depth = 1
  913.     if options.test_deep:
  914.         depth = 2
  915.  
  916.     kwargs = {}
  917.     if options.test_red:
  918.         kwargs["red"] = options.test_red
  919.     if options.test_green:
  920.         kwargs["green"] = options.test_green
  921.     if options.test_blue:
  922.         kwargs["blue"] = options.test_blue
  923.     if options.test_alpha:
  924.         kwargs["alpha"] = options.test_alpha
  925.     pixels = test_rgba(size, depth, **kwargs)
  926.  
  927.     writer = Writer(size, size,
  928.                     bytes_per_sample=depth,
  929.                     transparent=options.transparent,
  930.                     background=options.background,
  931.                     gamma=options.gamma,
  932.                     has_alpha=options.test_alpha,
  933.                     compression=options.compression,
  934.                     interlaced=options.interlace)
  935.     writer.write_array(sys.stdout, pixels)
  936.  
  937.  
  938. def read_pnm_header(infile, supported='P6'):
  939.     """
  940.     Read a PNM header, return width and height of the image in pixels.
  941.     """
  942.     header = []
  943.     while len(header) < 4:
  944.         line = infile.readline()
  945.         sharp = line.find('#')
  946.         if sharp > -1:
  947.             line = line[:sharp]
  948.         header.extend(line.split())
  949.         if len(header) == 3 and header[0] == 'P4':
  950.             break # PBM doesn't have maxval
  951.     if header[0] not in supported:
  952.         raise NotImplementedError('file format %s not supported' % header[0])
  953.     if header[0] != 'P4' and header[3] != '255':
  954.         raise NotImplementedError('maxval %s not supported' % header[3])
  955.     return int(header[1]), int(header[2])
  956.  
  957.  
  958. def color_triple(color):
  959.     """
  960.     Convert a command line color value to a RGB triple of integers.
  961.     FIXME: Somewhere we need support for greyscale backgrounds etc.
  962.     """
  963.     if color.startswith('#') and len(color) == 4:
  964.         return (int(color[1], 16),
  965.                 int(color[2], 16),
  966.                 int(color[3], 16))
  967.     if color.startswith('#') and len(color) == 7:
  968.         return (int(color[1:3], 16),
  969.                 int(color[3:5], 16),
  970.                 int(color[5:7], 16))
  971.     elif color.startswith('#') and len(color) == 13:
  972.         return (int(color[1:5], 16),
  973.                 int(color[5:9], 16),
  974.                 int(color[9:13], 16))
  975.  
  976.  
  977. def _main():
  978.     """
  979.     Run the PNG encoder with options from the command line.
  980.     """
  981.     # Parse command line arguments
  982.     from optparse import OptionParser
  983.     version = '%prog ' + __revision__.strip('$').replace('Rev: ', 'r')
  984.     parser = OptionParser(version=version)
  985.     parser.set_usage("%prog [options] [pnmfile]")
  986.     parser.add_option("-i", "--interlace",
  987.                       default=False, action="store_true",
  988.                       help="create an interlaced PNG file (Adam7)")
  989.     parser.add_option("-t", "--transparent",
  990.                       action="store", type="string", metavar="color",
  991.                       help="mark the specified color as transparent")
  992.     parser.add_option("-b", "--background",
  993.                       action="store", type="string", metavar="color",
  994.                       help="save the specified background color")
  995.     parser.add_option("-a", "--alpha",
  996.                       action="store", type="string", metavar="pgmfile",
  997.                       help="alpha channel transparency (RGBA)")
  998.     parser.add_option("-g", "--gamma",
  999.                       action="store", type="float", metavar="value",
  1000.                       help="save the specified gamma value")
  1001.     parser.add_option("-c", "--compression",
  1002.                       action="store", type="int", metavar="level",
  1003.                       help="zlib compression level (0-9)")
  1004.     parser.add_option("-T", "--test",
  1005.                       default=False, action="store_true",
  1006.                       help="create a test image")
  1007.     parser.add_option("-R", "--test-red",
  1008.                       action="store", type="string", metavar="pattern",
  1009.                       help="test pattern for the red image layer")
  1010.     parser.add_option("-G", "--test-green",
  1011.                       action="store", type="string", metavar="pattern",
  1012.                       help="test pattern for the green image layer")
  1013.     parser.add_option("-B", "--test-blue",
  1014.                       action="store", type="string", metavar="pattern",
  1015.                       help="test pattern for the blue image layer")
  1016.     parser.add_option("-A", "--test-alpha",
  1017.                       action="store", type="string", metavar="pattern",
  1018.                       help="test pattern for the alpha image layer")
  1019.     parser.add_option("-D", "--test-deep",
  1020.                       default=False, action="store_true",
  1021.                       help="use test patterns with 16 bits per layer")
  1022.     parser.add_option("-S", "--test-size",
  1023.                       action="store", type="int", metavar="size",
  1024.                       help="width and height of the test image")
  1025.     (options, args) = parser.parse_args()
  1026.  
  1027.     # Convert options
  1028.     if options.transparent is not None:
  1029.         options.transparent = color_triple(options.transparent)
  1030.     if options.background is not None:
  1031.         options.background = color_triple(options.background)
  1032.  
  1033.     # Run regression tests
  1034.     if options.test:
  1035.         return test_suite(options)
  1036.  
  1037.     # Prepare input and output files
  1038.     if len(args) == 0:
  1039.         ppmfilename = '-'
  1040.         ppmfile = sys.stdin
  1041.     elif len(args) == 1:
  1042.         ppmfilename = args[0]
  1043.         ppmfile = open(ppmfilename, 'rb')
  1044.     else:
  1045.         parser.error("more than one input file")
  1046.     outfile = sys.stdout
  1047.  
  1048.     # Encode PNM to PNG
  1049.     width, height = read_pnm_header(ppmfile)
  1050.     writer = Writer(width, height,
  1051.                     interlaced=options.interlace,
  1052.                     transparent=options.transparent,
  1053.                     background=options.background,
  1054.                     has_alpha=options.alpha is not None,
  1055.                     gamma=options.gamma,
  1056.                     compression=options.compression)
  1057.     if options.alpha is not None:
  1058.         pgmfile = open(options.alpha, 'rb')
  1059.         awidth, aheight = read_pnm_header(pgmfile, 'P5')
  1060.         if (awidth, aheight) != (width, height):
  1061.             raise ValueError("alpha channel image size mismatch" +
  1062.                              " (%s has %sx%s but %s has %sx%s)"
  1063.                              % (ppmfilename, width, height,
  1064.                                 options.alpha, awidth, aheight))
  1065.         writer.convert_ppm_and_pgm(ppmfile, pgmfile, outfile)
  1066.     else:
  1067.         writer.convert_ppm(ppmfile, outfile)
  1068.  
  1069.  
  1070. if __name__ == '__main__':
  1071.     _main()
  1072.