Monday, August 11, 2014

Writing PNG files in Python

Writing PNG files in Python

A while ago I was trying to work out how to show a chart of some values on a web based UI. I'd originally though it might be fun to try and write PNGs from scratch, especially as the filtering included lends itself well to blocks of solid color, exactly what you'd expect from a simple 2-color bar chart. I got discouraged when looking at the PNG spec, and all of the libraries wanted to deal with image data (I really just wanted bars of a certain length, and didn't necessarily want to create an image out of them), so I gave up and used SVG instead.

But I started looking into it again recently to see if it really was that complicated to write PNGs, and after working out that the built-in zlib module has the compression and CRC functionality already there, it turned out to be quite simple.

To write a PNG block, you specify the length of the data, a 4-byte header, the data itself, and a CRC of everything except the length:

def yield_block(header, data):
    assert len(header)==4, 'header must be 4 bytes!'
    # length:
    yield struct.pack('! L', len(data))
    # chunk type, 4 byte header
    yield header
    # data
    yield data
    # crc
    yield struct.pack(
        '! L',
        zlib.crc32("".join([header, data])) & 0xffffffff)

I decided to yield everything instead of writing to a file, because of something.

Then it's just a case of writing a magic header, a 'IHDR' block with information about the image, and 'IDAT' block with the data, and finally an empty 'IEND' block.

I decided to write a function that takes a list of numbers and outputs a simple 1-bit png. Each row has a number of white pixels as specified in the list, followed by enough black pixels to finish the line. I decided not to filter the scan lines in the end (I think it would make more sense to do so if I were to use bytes instead of bits for the pixel (and I didn't check to see if bits really saved space over bytes either)):

def make_bar_png(data, dmax='auto'):
    def make_line(line):
        r = [0xFF] * (line // 8)
        remaining_zeros = 8 - line % 8
        if remaining_zeros != 8:
            r.append(0xff ^ ((1 << remaining_zeros) - 1))
        r += [0x00] * (bytes_per_line - len(r))
        return r
    if dmax=='auto':
        dmax = max(data)
    bytes_per_line = (dmax+7) // 8
    width = dmax  
    height = len(data)
    bit_depth = 1
    color_type = 0  # grayscale
    compression_method, filter_method, interlace_method = 0, 0, 0
    yield bytearray([0x89, 'P', 'N', 'G', '\r', '\n', 0x1A, '\n'])
    # our header block
    for b in yield_block('IHDR', 
                         struct.pack('! LLBBBBB',
                                     width, height, bit_depth,
                                     color_type, compression_method, 
                                     filter_method, interlace_method)):
        yield b
    #unfiltered data (start with 0 as filterbyte)
    dat = [str(bytearray([0]+make_line(x))) for x in data]
    for b in yield_block('IDAT',
        yield b
    for b in yield_block('IEND', ''):
        yield b

And that's it, a simple horizontal PNG bar graph generator, just using the Python standard library.

Here's the first 12 digits of pi, as a (very small) bar graph, enjoy.

The code, and a couple of simple tests, are on github. It could probably be extended to make vertical bar graphs, and use the difference-of-preceding byte (or byte above) filtering to make the files smaller, but it works as it is.