Pure python PNG writer

make_png.py

import numpy as np
import zlib
import struct

def make_png(data):
    """
    Convert numpy array to PNG byte array.
    
    *data* must be (H, W, 4) with dtype=ubyte    
    """
    
    # www.libpng.org/pub/png/spec/1.2/PNG-Structure.html
    header = b'\x89PNG\x0d\x0a\x1a\x0a' # header

    def mkchunk(data, name):
        if isinstance(data, np.ndarray):
            size = data.nbytes
        else:
            size = len(data)
        chunk = np.empty(size + 12, dtype=np.ubyte)
        chunk.data[0:4] = struct.pack('!I', size)
        chunk.data[4:8] = name # b'CPXS' # critical, public, standard, safe
        chunk.data[8:8+size] = data
        chunk.data[-4:] = struct.pack('!i', zlib.crc32(chunk[4:-4]))
        return chunk

    # www.libpng.org/pub/png/spec/1.2/PNG-Chunks.html#C.IHDR
    ctyp = 0b0110  # alpha, color
    h, w = data.shape[:2]
    depth = data.itemsize * 8
    ihdr = struct.pack('!IIBBBBB', w, h, depth, ctyp, 0, 0, 0)
    c1 = mkchunk(ihdr, 'IHDR')

    # www.libpng.org/pub/png/spec/1.2/PNG-Chunks.html#C.IDAT
    idat = np.empty((h, w*4 + 1), dtype=np.ubyte) # insert filter byte at each scanline
    idat[:, 1:] = data.reshape(h, w*4)
    idat[:, 0] = 0
    c2 = mkchunk(zlib.compress(idat), 'IDAT')

    c3 = mkchunk(np.empty((0,), dtype=np.ubyte), 'IEND')

    # concatenate
    lh = len(header)
    png = np.empty(lh + c1.nbytes + c2.nbytes + c3.nbytes, dtype=np.ubyte)
    png.data[:lh] = header
    p = lh
    for chunk in (c1, c2, c3):
        png[p:p+len(chunk)] = chunk
        p += chunk.nbytes

    return png

if __name__ == '__main__':
    data = np.zeros((100,100,4), dtype=np.ubyte)
    data[:,:,3] = 255
    data[:,:,0] = 100
    data[10,:] = (0,0,0,255)
    data[:,10] = (0,0,0,255)
    open('test.png', 'w').write(make_png(data))