/* -*- mode: c++; c-basic-offset: 4 -*- */

// this code is heavily adapted from the paint license, which is in

// the file paint.license (BSD compatible) included in this

// distribution. TODO, add license file to MANIFEST.in and CVS

/* For linux, png.h must be imported before Python.h because

png.h needs to be the one to define setjmp.

Undefining _POSIX_C_SOURCE and _XOPEN_SOURCE stops a couple

of harmless warnings.

*/

extern "C" {

# include <png.h>

# ifdef _POSIX_C_SOURCE

# undef _POSIX_C_SOURCE

# endif

# ifdef _XOPEN_SOURCE

# undef _XOPEN_SOURCE

# endif

}

#include "numpy_cpp.h"

#include "mplutils.h"

#include "file_compat.h"

# include <vector>

# include "Python.h"

// As reported in [3082058] build _png.so on aix

#ifdef _AIX

#undef jmpbuf

#endif

struct buffer_t {

PyObject *str;

size_t cursor;

size_t size;

};

static void write_png_data_buffer(png_structp png_ptr, png_bytep data, png_size_t length)

{

buffer_t *buff = (buffer_t *)png_get_io_ptr(png_ptr);

if (buff->cursor + length < buff->size) {

memcpy(PyBytes_AS_STRING(buff->str) + buff->cursor, data, length);

buff->cursor += length;

}

}

static void flush_png_data_buffer(png_structp png_ptr)

{

}

static void write_png_data(png_structp png_ptr, png_bytep data, png_size_t length)

{

PyObject *py_file_obj = (PyObject *)png_get_io_ptr(png_ptr);

PyObject *write_method = PyObject_GetAttrString(py_file_obj, "write");

PyObject *result = NULL;

if (write_method) {

#if PY3K

result = PyObject_CallFunction(write_method, (char *)"y#", data, length);

#else

result = PyObject_CallFunction(write_method, (char *)"s#", data, length);

#endif

}

Py_XDECREF(write_method);

Py_XDECREF(result);

}

static void flush_png_data(png_structp png_ptr)

{

PyObject *py_file_obj = (PyObject *)png_get_io_ptr(png_ptr);

PyObject *flush_method = PyObject_GetAttrString(py_file_obj, "flush");

PyObject *result = NULL;

if (flush_method) {

result = PyObject_CallFunction(flush_method, (char *)"");

}

Py_XDECREF(flush_method);

Py_XDECREF(result);

}

const char *Py_write_png__doc__ =

"write_png(buffer, file, dpi=0, compression=6, filter=auto)\n"

"\n"

"Parameters\n"

"----------\n"

"buffer : numpy array of image data\n"

" Must be an MxNxD array of dtype uint8.\n"

" - if D is 1, the image is greyscale\n"

" - if D is 3, the image is RGB\n"

" - if D is 4, the image is RGBA\n"

"\n"

"file : str path, file-like object or None\n"

" - If a str, must be a file path\n"

" - If a file-like object, must write bytes\n"

" - If None, a byte string containing the PNG data will be returned\n"

"\n"

"dpi : float\n"

" The dpi to store in the file metadata.\n"

"\n"

"compression : int\n"

" The level of lossless zlib compression to apply. 0 indicates no\n"

" compression. Values 1-9 indicate low/fast through high/slow\n"

" compression. Default is 6.\n"

"\n"

"filter : int\n"

" Filter to apply. Must be one of the constants: PNG_FILTER_NONE,\n"

" PNG_FILTER_SUB, PNG_FILTER_UP, PNG_FILTER_AVG, PNG_FILTER_PAETH.\n"

" See the PNG standard for more information.\n"

" If not provided, libpng will try to automatically determine the\n"

" best filter on a line-by-line basis.\n"

"\n"

"Returns\n"

"-------\n"

"buffer : bytes or None\n"

" Byte string containing the PNG content if None was passed in for\n"

" file, otherwise None is returned.\n";

static PyObject *Py_write_png(PyObject *self, PyObject *args, PyObject *kwds)

{

numpy::array_view<unsigned char, 3> buffer;

PyObject *filein;

PyObject *metadata = NULL;

PyObject *meta_key, *meta_val;

png_text *text;

Py_ssize_t pos = 0;

int meta_pos = 0;

Py_ssize_t meta_size;

double dpi = 0;

int compression = 6;

int filter = -1;

const char *names[] = { "buffer", "file", "dpi", "compression", "filter", "metadata", NULL };

// We don't need strict contiguity, just for each row to be

// contiguous, and libpng has special handling for getting RGB out

// of RGBA, ARGB or BGR. But the simplest thing to do is to

// enforce contiguity using array_view::converter_contiguous.

if (!PyArg_ParseTupleAndKeywords(args,

kwds,

"O&O|diiO:write_png",

(char **)names,

&buffer.converter_contiguous,

&buffer,

&filein,

&dpi,

&compression,

&filter,

&metadata)) {

return NULL;

}

png_uint_32 width = (png_uint_32)buffer.dim(1);

png_uint_32 height = (png_uint_32)buffer.dim(0);

int channels = buffer.dim(2);

std::vector<png_bytep> row_pointers(height);

for (png_uint_32 row = 0; row < (png_uint_32)height; ++row) {

row_pointers[row] = (png_bytep)&buffer(row, 0, 0);

}

FILE *fp = NULL;

mpl_off_t offset = 0;

bool close_file = false;

bool close_dup_file = false;

PyObject *py_file = NULL;

png_structp png_ptr = NULL;

png_infop info_ptr = NULL;

struct png_color_8_struct sig_bit;

int png_color_type;

buffer_t buff;

buff.str = NULL;

switch (channels) {

case 1:

png_color_type = PNG_COLOR_TYPE_GRAY;

break;

case 3:

png_color_type = PNG_COLOR_TYPE_RGB;

break;

case 4:

png_color_type = PNG_COLOR_TYPE_RGB_ALPHA;

break;

default:

PyErr_SetString(PyExc_ValueError,

"Buffer must be an NxMxD array with D in 1, 3, 4 "

"(grayscale, RGB, RGBA)");

goto exit;

}

if (compression < 0 || compression > 9) {

PyErr_Format(PyExc_ValueError,

"compression must be in range 0-9, got %d", compression);

goto exit;

}

if (PyBytes_Check(filein) || PyUnicode_Check(filein)) {

if ((py_file = mpl_PyFile_OpenFile(filein, (char *)"wb")) == NULL) {

goto exit;

}

close_file = true;

} else {

py_file = filein;

}

if (filein == Py_None) {

buff.size = width * height * 4 + 1024;

buff.str = PyBytes_FromStringAndSize(NULL, buff.size);

if (buff.str == NULL) {

goto exit;

}

buff.cursor = 0;

} else {

#if PY3K

if (close_file) {

#else

if (close_file || PyFile_Check(py_file)) {

#endif

fp = mpl_PyFile_Dup(py_file, (char *)"wb", &offset);

}

if (fp) {

close_dup_file = true;

} else {

PyErr_Clear();

PyObject *write_method = PyObject_GetAttrString(py_file, "write");

if (!(write_method && PyCallable_Check(write_method))) {

Py_XDECREF(write_method);

PyErr_SetString(PyExc_TypeError,

"Object does not appear to be a 8-bit string path or "

"a Python file-like object");

goto exit;

}

Py_XDECREF(write_method);

}

}

png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);

if (png_ptr == NULL) {

PyErr_SetString(PyExc_RuntimeError, "Could not create write struct");

goto exit;

}

png_set_compression_level(png_ptr, compression);

if (filter >= 0) {

png_set_filter(png_ptr, 0, filter);

}

info_ptr = png_create_info_struct(png_ptr);

if (info_ptr == NULL) {

PyErr_SetString(PyExc_RuntimeError, "Could not create info struct");

goto exit;

}

if (setjmp(png_jmpbuf(png_ptr))) {

PyErr_SetString(PyExc_RuntimeError, "libpng signaled error");

goto exit;

}

if (buff.str) {

png_set_write_fn(png_ptr, (void *)&buff, &write_png_data_buffer, &flush_png_data_buffer);

} else if (fp) {

png_init_io(png_ptr, fp);

} else {

png_set_write_fn(png_ptr, (void *)py_file, &write_png_data, &flush_png_data);

}

png_set_IHDR(png_ptr,

info_ptr,

width,

height,

8,

png_color_type,

PNG_INTERLACE_NONE,

PNG_COMPRESSION_TYPE_BASE,

PNG_FILTER_TYPE_BASE);

// Save the dpi of the image in the file

if (dpi > 0.0) {

png_uint_32 dots_per_meter = (png_uint_32)(dpi / (2.54 / 100.0));

png_set_pHYs(png_ptr, info_ptr, dots_per_meter, dots_per_meter, PNG_RESOLUTION_METER);

}

#ifdef PNG_TEXT_SUPPORTED

// Save the metadata

if (metadata != NULL) {

meta_size = PyDict_Size(metadata);

text = new png_text[meta_size];

while (PyDict_Next(metadata, &pos, &meta_key, &meta_val)) {

text[meta_pos].compression = PNG_TEXT_COMPRESSION_NONE;

#if PY3K

if (PyUnicode_Check(meta_key)) {

PyObject *temp_key = PyUnicode_AsEncodedString(meta_key, "ASCII", "strict");

if (temp_key != NULL) {

text[meta_pos].key = PyBytes_AsString(temp_key);

}

} else if (PyBytes_Check(meta_key)) {

text[meta_pos].key = PyBytes_AsString(meta_key);

} else {

text[meta_pos].key = NULL; // Silently drops entry

}

if (PyUnicode_Check(meta_val)) {

PyObject *temp_val = PyUnicode_AsEncodedString(meta_val, "ASCII", "strict");

if (temp_val != NULL) {

text[meta_pos].text = PyBytes_AsString(temp_val);

}

} else if (PyBytes_Check(meta_val)) {

text[meta_pos].text = PyBytes_AsString(meta_val);

} else {

text[meta_pos].text = (char *)"Invalid value in metadata";

}

#else

text[meta_pos].key = PyString_AsString(meta_key);

text[meta_pos].text = PyString_AsString(meta_val);

#endif

#ifdef PNG_iTXt_SUPPORTED

text[meta_pos].lang = NULL;

#endif

meta_pos++;

}

png_set_text(png_ptr, info_ptr, text, meta_size);

delete[] text;

}

#endif

sig_bit.alpha = 0;

switch (png_color_type) {

case PNG_COLOR_TYPE_GRAY:

sig_bit.gray = 8;

sig_bit.red = 0;

sig_bit.green = 0;

sig_bit.blue = 0;

break;

case PNG_COLOR_TYPE_RGB_ALPHA:

sig_bit.alpha = 8;

// fall through

case PNG_COLOR_TYPE_RGB:

sig_bit.gray = 0;

sig_bit.red = 8;

sig_bit.green = 8;

sig_bit.blue = 8;

break;

default:

PyErr_SetString(PyExc_RuntimeError, "internal error, bad png_color_type");

goto exit;

}

png_set_sBIT(png_ptr, info_ptr, &sig_bit);

png_write_info(png_ptr, info_ptr);

png_write_image(png_ptr, &row_pointers[0]);

png_write_end(png_ptr, info_ptr);

exit:

if (png_ptr && info_ptr) {

png_destroy_write_struct(&png_ptr, &info_ptr);

}

if (close_dup_file) {

mpl_PyFile_DupClose(py_file, fp, offset);

}

if (close_file) {

mpl_PyFile_CloseFile(py_file);

Py_DECREF(py_file);

}

if (PyErr_Occurred()) {

Py_XDECREF(buff.str);

return NULL;

} else {

if (buff.str) {

_PyBytes_Resize(&buff.str, buff.cursor);

return buff.str;

}

Py_RETURN_NONE;

}

}

static void _read_png_data(PyObject *py_file_obj, png_bytep data, png_size_t length)

{

PyObject *read_method = PyObject_GetAttrString(py_file_obj, "read");

PyObject *result = NULL;

char *buffer;

Py_ssize_t bufflen;

if (read_method) {

result = PyObject_CallFunction(read_method, (char *)"i", length);

if (PyBytes_AsStringAndSize(result, &buffer, &bufflen) == 0) {

if (bufflen == (Py_ssize_t)length) {

memcpy(data, buffer, length);

} else {

PyErr_SetString(PyExc_IOError, "read past end of file");

}

}

}

Py_XDECREF(read_method);

Py_XDECREF(result);

}

static void read_png_data(png_structp png_ptr, png_bytep data, png_size_t length)

{

PyObject *py_file_obj = (PyObject *)png_get_io_ptr(png_ptr);

_read_png_data(py_file_obj, data, length);

}

static PyObject *_read_png(PyObject *filein, bool float_result)

{

png_byte header[8]; // 8 is the maximum size that can be checked

FILE *fp = NULL;

mpl_off_t offset = 0;

bool close_file = false;

bool close_dup_file = false;

PyObject *py_file = NULL;

png_structp png_ptr = NULL;

png_infop info_ptr = NULL;

int num_dims;

std::vector<png_bytep> row_pointers;

png_uint_32 width = 0;

png_uint_32 height = 0;

int bit_depth;

PyObject *result = NULL;

// TODO: Remove direct calls to Numpy API here

if (PyBytes_Check(filein) || PyUnicode_Check(filein)) {

if ((py_file = mpl_PyFile_OpenFile(filein, (char *)"rb")) == NULL) {

goto exit;

}

close_file = true;

} else {

py_file = filein;

}

#if PY3K

if (close_file) {

#else

if (close_file || PyFile_Check(py_file)) {

#endif

fp = mpl_PyFile_Dup(py_file, (char *)"rb", &offset);

}

if (fp) {

close_dup_file = true;

if (fread(header, 1, 8, fp) != 8) {

PyErr_SetString(PyExc_IOError, "error reading PNG header");

goto exit;

}

} else {

PyErr_Clear();

PyObject *read_method = PyObject_GetAttrString(py_file, "read");

if (!(read_method && PyCallable_Check(read_method))) {

Py_XDECREF(read_method);

PyErr_SetString(PyExc_TypeError,

"Object does not appear to be a 8-bit string path or "

"a Python file-like object");

goto exit;

}

Py_XDECREF(read_method);

_read_png_data(py_file, header, 8);

}

if (png_sig_cmp(header, 0, 8)) {

PyErr_SetString(PyExc_ValueError, "invalid PNG header");

goto exit;

}

/* initialize stuff */

png_ptr = png_create_read_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);

if (!png_ptr) {

PyErr_SetString(PyExc_RuntimeError, "png_create_read_struct failed");

goto exit;

}

info_ptr = png_create_info_struct(png_ptr);

if (!info_ptr) {

PyErr_SetString(PyExc_RuntimeError, "png_create_info_struct failed");

goto exit;

}

if (setjmp(png_jmpbuf(png_ptr))) {

PyErr_SetString(PyExc_RuntimeError, "Error setting jump");

goto exit;

}

if (fp) {

png_init_io(png_ptr, fp);

} else {

png_set_read_fn(png_ptr, (void *)py_file, &read_png_data);

}

png_set_sig_bytes(png_ptr, 8);

png_read_info(png_ptr, info_ptr);

width = png_get_image_width(png_ptr, info_ptr);

height = png_get_image_height(png_ptr, info_ptr);

bit_depth = png_get_bit_depth(png_ptr, info_ptr);

// Unpack 1, 2, and 4-bit images

if (bit_depth < 8) {

png_set_packing(png_ptr);

}

// If sig bits are set, shift data

png_color_8p sig_bit;

if ((png_get_color_type(png_ptr, info_ptr) != PNG_COLOR_TYPE_PALETTE) &&

png_get_sBIT(png_ptr, info_ptr, &sig_bit)) {

png_set_shift(png_ptr, sig_bit);

}

// Convert big endian to little

if (bit_depth == 16) {

png_set_swap(png_ptr);

}

// Convert palletes to full RGB

if (png_get_color_type(png_ptr, info_ptr) == PNG_COLOR_TYPE_PALETTE) {

png_set_palette_to_rgb(png_ptr);

bit_depth = 8;

}

// If there's an alpha channel convert gray to RGB

if (png_get_color_type(png_ptr, info_ptr) == PNG_COLOR_TYPE_GRAY_ALPHA) {

png_set_gray_to_rgb(png_ptr);

}

png_set_interlace_handling(png_ptr);

png_read_update_info(png_ptr, info_ptr);

row_pointers.resize(height);

for (png_uint_32 row = 0; row < height; row++) {

row_pointers[row] = new png_byte[png_get_rowbytes(png_ptr, info_ptr)];

}

png_read_image(png_ptr, &row_pointers[0]);

npy_intp dimensions[3];

dimensions[0] = height; // numrows

dimensions[1] = width; // numcols

if (png_get_color_type(png_ptr, info_ptr) & PNG_COLOR_MASK_ALPHA) {

dimensions[2] = 4; // RGBA images

} else if (png_get_color_type(png_ptr, info_ptr) & PNG_COLOR_MASK_COLOR) {

dimensions[2] = 3; // RGB images

} else {

dimensions[2] = 1; // Greyscale images

}

if (float_result) {

double max_value = (1 << bit_depth) - 1;

numpy::array_view<float, 3> A(dimensions);

for (png_uint_32 y = 0; y < height; y++) {

png_byte *row = row_pointers[y];

for (png_uint_32 x = 0; x < width; x++) {

if (bit_depth == 16) {

png_uint_16 *ptr = &reinterpret_cast<png_uint_16 *>(row)[x * dimensions[2]];

for (png_uint_32 p = 0; p < (png_uint_32)dimensions[2]; p++) {

A(y, x, p) = (float)(ptr[p]) / max_value;

}

} else {

png_byte *ptr = &(row[x * dimensions[2]]);

for (png_uint_32 p = 0; p < (png_uint_32)dimensions[2]; p++) {

A(y, x, p) = (float)(ptr[p]) / max_value;

}

}

}

}

result = A.pyobj();

} else if (bit_depth == 16) {

numpy::array_view<png_uint_16, 3> A(dimensions);

for (png_uint_32 y = 0; y < height; y++) {

png_byte *row = row_pointers[y];

for (png_uint_32 x = 0; x < width; x++) {

png_uint_16 *ptr = &reinterpret_cast<png_uint_16 *>(row)[x * dimensions[2]];

for (png_uint_32 p = 0; p < (png_uint_32)dimensions[2]; p++) {

A(y, x, p) = ptr[p];

}

}

}

result = A.pyobj();

} else if (bit_depth == 8) {

numpy::array_view<png_byte, 3> A(dimensions);

for (png_uint_32 y = 0; y < height; y++) {

png_byte *row = row_pointers[y];

for (png_uint_32 x = 0; x < width; x++) {

png_byte *ptr = &(row[x * dimensions[2]]);

for (png_uint_32 p = 0; p < (png_uint_32)dimensions[2]; p++) {

A(y, x, p) = ptr[p];

}

}

}

result = A.pyobj();

} else {

PyErr_SetString(PyExc_RuntimeError, "image has unknown bit depth");

goto exit;

}

// free the png memory

png_read_end(png_ptr, info_ptr);

// For gray, return an x by y array, not an x by y by 1

num_dims = (png_get_color_type(png_ptr, info_ptr) & PNG_COLOR_MASK_COLOR) ? 3 : 2;

if (num_dims == 2) {

PyArray_Dims dims = {dimensions, 2};

PyObject *reshaped = PyArray_Newshape((PyArrayObject *)result, &dims, NPY_CORDER);

Py_DECREF(result);

result = reshaped;

}

exit:

if (png_ptr && info_ptr) {

#ifndef png_infopp_NULL

png_destroy_read_struct(&png_ptr, &info_ptr, NULL);

#else

png_destroy_read_struct(&png_ptr, &info_ptr, png_infopp_NULL);

#endif

}

if (close_dup_file) {

mpl_PyFile_DupClose(py_file, fp, offset);

}

if (close_file) {

mpl_PyFile_CloseFile(py_file);

Py_DECREF(py_file);

}

for (png_uint_32 row = 0; row < height; row++) {

delete[] row_pointers[row];

}

if (PyErr_Occurred()) {

Py_XDECREF(result);

return NULL;

} else {

return result;

}

}

const char *Py_read_png_float__doc__ =

"read_png_float(file)\n"

"\n"

"Read in a PNG file, converting values to floating-point doubles\n"

"in the range (0, 1)\n"

"\n"

"Parameters\n"

"----------\n"

"file : str path or file-like object\n";

static PyObject *Py_read_png_float(PyObject *self, PyObject *args, PyObject *kwds)

{

return _read_png(args, true);

}

const char *Py_read_png_int__doc__ =

"read_png_int(file)\n"

"\n"

"Read in a PNG file with original integer values.\n"

"\n"

"Parameters\n"

"----------\n"

"file : str path or file-like object\n";

static PyObject *Py_read_png_int(PyObject *self, PyObject *args, PyObject *kwds)

{

return _read_png(args, false);

}

const char *Py_read_png__doc__ =

"read_png(file)\n"

"\n"

"Read in a PNG file, converting values to floating-point doubles\n"

"in the range (0, 1)\n"

"\n"

"Alias for read_png_float()\n"

"\n"

"Parameters\n"

"----------\n"

"file : str path or file-like object\n";

static PyMethodDef module_methods[] = {

{"write_png", (PyCFunction)Py_write_png, METH_VARARGS|METH_KEYWORDS, Py_write_png__doc__},

{"read_png", (PyCFunction)Py_read_png_float, METH_O, Py_read_png__doc__},

{"read_png_float", (PyCFunction)Py_read_png_float, METH_O, Py_read_png_float__doc__},

{"read_png_int", (PyCFunction)Py_read_png_int, METH_O, Py_read_png_int__doc__},

{NULL}

};

extern "C" {

#if PY3K

static struct PyModuleDef moduledef = {

PyModuleDef_HEAD_INIT,

"_png",

NULL,

0,

module_methods,

NULL,

NULL,

NULL,

NULL

};

#define INITERROR return NULL

PyMODINIT_FUNC PyInit__png(void)

#else

#define INITERROR return

PyMODINIT_FUNC init_png(void)

#endif

{

PyObject *m;

#if PY3K

m = PyModule_Create(&moduledef);

#else

m = Py_InitModule3("_png", module_methods, NULL);

#endif

if (m == NULL) {

INITERROR;

}

import_array();

if (PyModule_AddIntConstant(m, "PNG_FILTER_NONE", PNG_FILTER_NONE) ||

PyModule_AddIntConstant(m, "PNG_FILTER_SUB", PNG_FILTER_SUB) ||

PyModule_AddIntConstant(m, "PNG_FILTER_UP", PNG_FILTER_UP) ||

PyModule_AddIntConstant(m, "PNG_FILTER_AVG", PNG_FILTER_AVG) ||

PyModule_AddIntConstant(m, "PNG_FILTER_PAETH", PNG_FILTER_PAETH)) {

INITERROR;

}

#if PY3K

return m;

#endif

}

}