[Numpy-discussion] saving and loading as PNG

[Martin Teichmann]

Hi all,

i am working with numpy and mostly 16 bit matrices, and so I was looking
for a standard way of saving them. I found that PNG actually supports
that, but I could not find a way to use this feature from within
python. I thought that it actually is a cool way of storing matrices
and thus I wrote a little pyrex program that does the job. I figured
this might be interesting to a broader audience, so I posted it here,
maybe for inclusion in numpy?

Some more thoughts:
* Other implementations: There is other people who have done such a thing.
the PIL knows how to read and write PNG, but only 8 bit. The same holds
for matplotlib.

* Problems: until now, numpy does not link agains libpng. This should not
be a problem on linux-like systems, as those normally all have libpng 
installed. Windows however, I don't know. Same thing with pyrex. 
Numpy doesn't use pyrex until now.

* other places to put: numpy could be considered to basic to have such
a functionality, one could put it in scipy. Well, but png is not really
a _scientific_ application (even though I personally use it to do science).
Matplotlib would be another nice place to put this feature, big advantage
is that matplotlib already links against libpng.

So, now I shared my thoughts with you, comments?

Greetings,

Martin

And here the code (188 lines are hopefully ok for a mailinglist?)

cimport c_numpy
import numpy

cdef extern from "stdlib.h":
    cdef void *malloc(int)
    cdef void free(void *)

cdef extern from "stdio.h":
    cdef struct file:
        pass
    ctypedef file FILE
    cdef FILE *fopen(char *, char *)
    cdef int fclose(FILE *)

cdef extern from "errno.h":
    cdef int errno

cdef extern from "png.h":
    cdef struct jmp_buf:
        pass
    cdef int setjmp(jmp_buf)
    cdef struct png_struct:
        jmp_buf jmpbuf
    ctypedef png_struct *png_structp

    cdef struct png_info:
        int width
        int height
        int bit_depth
        int color_type
        int channels
    ctypedef png_info *png_infop

    cdef struct png_text:
        int compression
        char *key
        char *text
        int text_length
    ctypedef png_text *png_textp

    cdef png_structp png_create_read_struct(char *, void *, void *, void*)
    cdef void png_init_io(png_structp, FILE *)
    cdef void png_set_sig_bytes(png_structp, int)
    cdef png_infop png_create_info_struct(png_structp)
    cdef void png_read_info(png_structp, png_infop)
    cdef int png_set_interlace_handling(png_structp)
    cdef void png_read_update_info(png_structp, png_infop)
    cdef void png_read_image(png_structp, unsigned char **)
    cdef void png_read_png(png_structp, png_infop, int, void *)
    cdef void png_read_end(png_structp, png_infop)
    cdef void png_destroy_read_struct(png_structp *, png_infop *, void *)
    cdef png_structp png_create_write_struct(char *, void *, void *, void*)
    cdef void png_set_IHDR(png_structp, png_infop, unsigned int,
        unsigned int, int, int, int, int, int)
    cdef void png_write_info(png_structp, png_infop)
    cdef void png_write_image(png_structp, unsigned char **)
    cdef void png_write_end(png_structp, png_infop)
    cdef void png_destroy_write_struct(png_structp *,png_infop *)
    cdef void png_set_swap(png_structp)
    cdef void png_set_rows(png_structp, png_infop, unsigned char **)
    cdef void png_write_png(png_structp, png_infop, int, void *)
    cdef void png_set_text(png_structp, png_infop, png_textp, int)
    cdef int PNG_COLOR_TYPE_GRAY
    cdef int PNG_COLOR_TYPE_PALETTE
    cdef int PNG_COLOR_TYPE_RGB
    cdef int PNG_COLOR_TYPE_RGB_ALPHA
    cdef int PNG_COLOR_TYPE_GRAY_ALPHA
    cdef int PNG_INTERLACE_NONE
    cdef int PNG_COMPRESSION_TYPE_BASE
    cdef int PNG_FILTER_TYPE_BASE
    cdef int PNG_TRANSFORM_IDENTITY
    cdef int PNG_TRANSFORM_SWAP_ENDIAN
    cdef char *PNG_LIBPNG_VER_STRING

def load(fn):
    """ load a PNG file into a NumPy array.

    This function can read all kinds of PNG images in a senseful way:
    Depending on the number of bits per pixel, we chose numpy.bool,
    numpy.unit8 or numpy.uint16 as base type. For gray scale images or
    pallettes, this returns a 2-D array, for everything else a 3-D array,
    where the third dimension has a size of 2 for gray+alpha images, 3 for
    RGB and 4 for RGBA images.  """

    cdef FILE *f
    f = fopen(fn, "rb")
    if f == NULL:
        raise IOError(errno)
    cdef png_structp p
    p = png_create_read_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL)
    if setjmp(p.jmpbuf):
        fclose(f)
        raise RuntimeError("libpng reported error")
    cdef png_infop i
    i = png_create_info_struct(p)
    png_init_io(p, f)
    png_read_info(p, i)

    if i.color_type == PNG_COLOR_TYPE_GRAY:
        shape = (i.height, i.width)
    elif i.color_type == PNG_COLOR_TYPE_GRAY_ALPHA:
        shape = (i.height, i.width, 2)
    elif i.color_type == PNG_COLOR_TYPE_PALETTE:
        shape = (i.height, i.width)
    elif i.color_type ==  PNG_COLOR_TYPE_RGB:
        shape = (i.height, i.width, 3)
    elif i.color_type == PNG_COLOR_TYPE_RGB_ALPHA:
        shape = (i.height, i.width, 4)

    if i.bit_depth == 1: dtype = numpy.bool
    elif i.bit_depth == 8: dtype = numpy.uint8
    elif i.bit_depth == 16: dtype = numpy.uint16

    png_set_interlace_handling(p)
    png_read_update_info(p, i)

    r = numpy.zeros(shape, dtype)
    cdef c_numpy.ndarray s
    s = r
    cdef unsigned char **rp
    rp = <unsigned char **> malloc(sizeof(char *) * i.height)
    cdef int j
    for j from 0 <= j < i.height:
        rp[j] = <unsigned char *> &s.data[(j *
            i.channels * i.width * i.bit_depth) / 8]
    png_set_swap(p)
    png_read_image(p, rp)
    free(rp)
    png_destroy_read_struct(&p, &i, NULL)
    fclose(f)
    return r

def save(fn, data):
    """ Save a NumPy array to PNG

    This function saves the NumPy array in data to the file fn. It saves
    the data the best way it can: 2-D arrays are saved as gray scale
    images, 3-D arrays are saved as gray+alpha, rgb or rgba depending on
    the size of the 3rd dimension.  """

    cdef c_numpy.ndarray d
    d = numpy.ascontiguousarray(data)
    cdef int j
    cdef int width
    if data.itemsize > 2:
        raise ValueError("PNG cannot handle BPP > 16")
    if data.ndim == 2:
        width = data.shape[1] * data.itemsize
    elif data.ndim == 3:
        width = data.shape[2] * data.shape[1] * data.itemsize
    else:
        raise ValueError("Can only save 2-D or 3-D arrays to PNG")
    cdef FILE *f
    f = fopen(fn, "wb")
    cdef png_structp p
    p = png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL)

    cdef unsigned char **rp
    rp = NULL
    if setjmp(p.jmpbuf):
        free(rp)
        fclose(f)
        raise RuntimeError("libpng reported error")
    cdef png_infop i
    i = png_create_info_struct(p)
    png_init_io(p, f)
    if data.ndim == 2:
        color = PNG_COLOR_TYPE_GRAY
    elif data.shape[2] == 2:
        color = PNG_COLOR_TYPE_GRAY_ALPHA
    elif data.shape[2] == 3:
        color = PNG_COLOR_TYPE_RGB
    elif data.shape[2] == 4:
        color = PNG_COLOR_TYPE_RGB_ALPHA
    else:
        raise ValueError("Third dimension must be <= 4 for saving to PNG")
    png_set_IHDR(p, i, data.shape[1], data.shape[0], data.itemsize*8,
        color, PNG_INTERLACE_NONE, PNG_COMPRESSION_TYPE_BASE,
        PNG_FILTER_TYPE_BASE)

    rp = <unsigned char **> malloc(sizeof(char *) * data.shape[0])
    for j from 0 <= j < data.shape[0]:
        rp[j] = <unsigned char *> &d.data[j * width]
    png_set_rows(p, i, rp)
    png_write_png(p, i, PNG_TRANSFORM_SWAP_ENDIAN, NULL)
    free(rp)
    png_destroy_write_struct(&p, &i)
    fclose(f)