// 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
Py::Object _png_module::write_png(const Py::Tuple& args)
{
args.verify_length(4, 5);
FILE *fp = NULL;
bool close_file = false;
Py::Object buffer_obj = Py::Object(args[0]);
PyObject* buffer = buffer_obj.ptr();
if (!PyObject_CheckReadBuffer(buffer)) {
throw Py::TypeError("First argument must be an rgba buffer.");
}
const void* pixBufferPtr = NULL;
Py_ssize_t pixBufferLength = 0;
if (PyObject_AsReadBuffer(buffer, &pixBufferPtr, &pixBufferLength)) {
throw Py::ValueError("Couldn't get data from read buffer.");
}
png_byte* pixBuffer = (png_byte*)pixBufferPtr;
int width = (int)Py::Int(args[1]);
int height = (int)Py::Int(args[2]);
if (pixBufferLength < width * height * 4) {
throw Py::ValueError("Buffer and width, height don't seem to match.");
}
Py::Object py_fileobj = Py::Object(args[3]);
if (py_fileobj.isString()) {
std::string fileName = Py::String(py_fileobj);
const char *file_name = fileName.c_str();
if ((fp = fopen(file_name, "wb")) == NULL)
throw Py::RuntimeError( Printf("Could not open file %s", file_name).str() );
close_file = true;
} else if (PyFile_CheckExact(py_fileobj.ptr())) {
fp = PyFile_AsFile(py_fileobj.ptr());
}
else {
PyObject* write_method = PyObject_GetAttrString(py_fileobj.ptr(), "write");
if (!(write_method && PyCallable_Check(write_method))) {
Py_XDECREF(write_method);
throw Py::TypeError("Object does not appear to be a 8-bit string path or a Python file-like object");
}
Py_XDECREF(write_method);
}
png_bytep *row_pointers = NULL;
png_structp png_ptr = NULL;
png_infop info_ptr = NULL;
try {
struct png_color_8_struct sig_bit;
png_uint_32 row;
row_pointers = new png_bytep[height];
for (row = 0; row < (png_uint_32)height; ++row) {
row_pointers[row] = pixBuffer + row * width * 4;
}
png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);
if (png_ptr == NULL) {
throw Py::RuntimeError("Could not create write struct");
}
info_ptr = png_create_info_struct(png_ptr);
if (info_ptr == NULL) {
throw Py::RuntimeError("Could not create info struct");
}
if (setjmp(png_ptr->jmpbuf)) {
throw Py::RuntimeError("Error building image");
}
if (fp) {
png_init_io(png_ptr, fp);
} else {
png_set_write_fn(png_ptr, (void*)py_fileobj.ptr(),
&write_png_data, &flush_png_data);
}
png_set_IHDR(png_ptr, info_ptr,
width, height, 8,
PNG_COLOR_TYPE_RGB_ALPHA, PNG_INTERLACE_NONE,
PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);
// Save the dpi of the image in the file
if (args.size() == 5) {
double dpi = Py::Float(args[4]);
size_t dots_per_meter = (size_t)(dpi / (2.54 / 100.0));
png_set_pHYs(png_ptr, info_ptr, dots_per_meter, dots_per_meter, PNG_RESOLUTION_METER);
}
// this a a color image!
sig_bit.gray = 0;
sig_bit.red = 8;
sig_bit.green = 8;
sig_bit.blue = 8;
/* if the image has an alpha channel then */
sig_bit.alpha = 8;
png_set_sBIT(png_ptr, info_ptr, &sig_bit);
png_write_info(png_ptr, info_ptr);
png_write_image(png_ptr, row_pointers);
png_write_end(png_ptr, info_ptr);
} catch (...) {
if (fp && close_file) fclose(fp);
delete [] row_pointers;
/* Changed calls to png_destroy_write_struct to follow
http://www.libpng.org/pub/png/libpng-manual.txt.
This ensures the info_ptr memory is released.
*/
if (png_ptr && info_ptr) png_destroy_write_struct(&png_ptr, &info_ptr);
throw;
}
png_destroy_write_struct(&png_ptr, &info_ptr);
delete [] row_pointers;
if (fp && close_file) fclose(fp);
return Py::Object();
}
Py::Object
_png_module::read_png(const Py::Tuple& args)
{
args.verify_length(1);
png_byte header[8]; // 8 is the maximum size that can be checked
FILE* fp = NULL;
bool close_file = false;
Py::Object py_fileobj = Py::Object(args[0]);
#if PY_MAJOR_VERSION >= 3
int fd = PyObject_AsFileDescriptor(py_fileobj.ptr());
PyErr_Clear();
#endif
if (py_fileobj.isString())
{
std::string fileName = Py::String(py_fileobj);
const char *file_name = fileName.c_str();
if ((fp = fopen(file_name, "rb")) == NULL)
{
throw Py::RuntimeError(
Printf("Could not open file %s for reading", file_name).str());
}
close_file = true;
}
#if PY_MAJOR_VERSION >= 3
else if (fd != -1) {
fp = fdopen(fd, "r");
}
#else
else if (PyFile_CheckExact(py_fileobj.ptr()))
{
fp = PyFile_AsFile(py_fileobj.ptr());
}
#endif
else
{
PyObject* read_method = PyObject_GetAttrString(py_fileobj.ptr(), "read");
if (!(read_method && PyCallable_Check(read_method)))
{
Py_XDECREF(read_method);
throw Py::TypeError("Object does not appear to be a 8-bit string path or a Python file-like object");
}
Py_XDECREF(read_method);
}
if (fp)
{
if (fread(header, 1, 8, fp) != 8)
{
throw Py::RuntimeError(
"_image_module::readpng: error reading PNG header");
}
}
else
{
_read_png_data(py_fileobj.ptr(), header, 8);
}
if (png_sig_cmp(header, 0, 8))
{
throw Py::RuntimeError(
"_image_module::readpng: file not recognized as a PNG file");
}
/* initialize stuff */
png_structp png_ptr = png_create_read_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);
if (!png_ptr)
{
throw Py::RuntimeError(
"_image_module::readpng: png_create_read_struct failed");
}
png_infop info_ptr = png_create_info_struct(png_ptr);
if (!info_ptr)
{
throw Py::RuntimeError(
"_image_module::readpng: png_create_info_struct failed");
}
if (setjmp(png_jmpbuf(png_ptr)))
{
throw Py::RuntimeError(
"_image_module::readpng: error during init_io");
}
if (fp)
{
png_init_io(png_ptr, fp);
}
else
{
png_set_read_fn(png_ptr, (void*)py_fileobj.ptr(), &read_png_data);
}
png_set_sig_bytes(png_ptr, 8);
png_read_info(png_ptr, info_ptr);
png_uint_32 width = info_ptr->width;
png_uint_32 height = info_ptr->height;
int bit_depth = info_ptr->bit_depth;
// 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 ((info_ptr->color_type != 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 (info_ptr->color_type == PNG_COLOR_TYPE_PALETTE)
{
png_set_palette_to_rgb(png_ptr);
}
// If there's an alpha channel convert gray to RGB
if (info_ptr->color_type == 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);
/* read file */
if (setjmp(png_jmpbuf(png_ptr)))
{
throw Py::RuntimeError(
"_image_module::readpng: error during read_image");
}
png_bytep *row_pointers = new png_bytep[height];
png_uint_32 row;
for (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);
npy_intp dimensions[3];
dimensions[0] = height; //numrows
dimensions[1] = width; //numcols
if (info_ptr->color_type & PNG_COLOR_MASK_ALPHA)
{
dimensions[2] = 4; //RGBA images
}
else if (info_ptr->color_type & PNG_COLOR_MASK_COLOR)
{
dimensions[2] = 3; //RGB images
}
else
{
dimensions[2] = 1; //Greyscale images
}
//For gray, return an x by y array, not an x by y by 1
int num_dims = (info_ptr->color_type & PNG_COLOR_MASK_COLOR) ? 3 : 2;
double max_value = (1 << ((bit_depth < 8) ? 8 : bit_depth)) - 1;
PyArrayObject *A = (PyArrayObject *) PyArray_SimpleNew(
num_dims, dimensions, PyArray_FLOAT);
if (A == NULL)
{
throw Py::MemoryError("Could not allocate image array");
}
for (png_uint_32 y = 0; y < height; y++)
{
png_byte* row = row_pointers[y];
for (png_uint_32 x = 0; x < width; x++)
{
size_t offset = y * A->strides[0] + x * A->strides[1];
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++)
{
*(float*)(A->data + offset + p*A->strides[2]) = (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++)
{
*(float*)(A->data + offset + p*A->strides[2]) = (float)(ptr[p]) / max_value;
}
}
}
}
//free the png memory
png_read_end(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_file)
{
fclose(fp);
}
for (row = 0; row < height; row++)
{
delete [] row_pointers[row];
}
delete [] row_pointers;
return Py::asObject((PyObject*)A);
}
static PyObject *
__write_dcd_header(PyObject *self, PyObject *args)
{
/* At this point we assume the file has been opened for writing */
PyObject *temp = NULL;
dcdhandle *dcd = NULL;
fio_fd fd;
int rc = 0;
int natoms = 0;
const char *remarks = "DCD";
int istart = 0; /* starting timestep of DCD file */
int nsavc = 1; /* number of timesteps between written DCD frames */
double delta = 1.0; /* length of a timestep */
int with_unitcell = 1; /* contains unit cell information (charmm format) */
int charmm = DCD_IS_CHARMM; /* charmm-formatted DCD file */
if (with_unitcell)
charmm |= DCD_HAS_EXTRA_BLOCK;
if (! self) {
/* we were in fact called as a module function, try to retrieve
a matching object from args */
if( !PyArg_ParseTuple(args, "Oi|iids", &self, &natoms, &istart, &nsavc, &delta, &remarks) )
return NULL;
} else {
/* we were obviously called as an object method so args should
only have the int value. */
if( !PyArg_ParseTuple(args, "i|iids", &natoms, &istart, &nsavc, &delta, &remarks) )
return NULL;
}
// Get the file object from the class
if (!PyObject_HasAttrString(self, "dcdfile")) {
// Raise exception
PyErr_SetString(PyExc_AttributeError, "dcdfile is not an attribute");
return NULL;
}
if ((temp = PyObject_GetAttrString(self, "dcdfile")) == NULL) { // This gives me a New Reference
// Raise exception
PyErr_SetString(PyExc_AttributeError, "dcdfile is not an attribute");
return NULL;
}
if (!PyFile_CheckExact(temp)) {
// Raise exception
PyErr_SetString(PyExc_TypeError, "dcdfile does not refer to a file object");
Py_DECREF(temp);
return NULL;
}
fd = fileno(PyFile_AsFile(temp));
// No longer need the reference to temp
Py_DECREF(temp);
dcd = (dcdhandle *)malloc(sizeof(dcdhandle));
memset(dcd, 0, sizeof(dcdhandle));
dcd->fd = fd;
if ((rc = write_dcdheader(dcd->fd, remarks, natoms, istart, nsavc, delta, with_unitcell, charmm)) < 0) {
PyErr_SetString(PyExc_IOError, "Cannot write header of DCD file");
free(dcd);
return NULL;
}
dcd->natoms = natoms;
dcd->nsets = 0;
dcd->istart = istart;
dcd->nsavc = nsavc;
dcd->delta = delta;
dcd->with_unitcell = with_unitcell;
dcd->charmm = charmm;
temp = PyCObject_FromVoidPtr(dcd, free); // Creates a New Reference
if (PyObject_SetAttrString(self, "_dcd_C_ptr", temp) == -1) {
// Raise exception - who knows what exception to raise??
PyErr_SetString(PyExc_AttributeError, "Could not create attribute _dcd_C_ptr");
Py_DECREF(temp);
return NULL;
}
Py_DECREF(temp);
// For debugging purposes
temp = PyBuffer_FromMemory(dcd, sizeof(dcdhandle)); // Creates a New Reference
if (PyObject_SetAttrString(self, "_dcd_C_str", temp) == -1) {
// Raise exception - who knows what exception to raise??
PyErr_SetString(PyExc_AttributeError, "Could not create attribute _dcd_C_str");
Py_DECREF(temp);
return NULL;
}
Py_DECREF(temp);
Py_INCREF(Py_None);
return Py_None;
}
PyObject*
verifymodule_ogg(PyObject *dummy, PyObject *args) {
PyObject *file_obj;
BitstreamReader *bitstream;
int has_previous_header = 0;
struct ogg_header previous_header;
struct ogg_header header;
uint8_t *data_buffer = NULL;
int data_buffer_size = 0;
int i;
uint32_t checksum;
/*fixes a "may be used unitialized" warning*/
previous_header.bitstream_serial_number =
previous_header.page_sequence_number = 0;
if (!PyArg_ParseTuple(args, "O", &file_obj))
return NULL;
if (!PyFile_CheckExact(file_obj)) {
PyErr_SetString(PyExc_TypeError,
"first argument must be an actual file object");
return NULL;
} else {
bitstream = br_open(PyFile_AsFile(file_obj), BS_LITTLE_ENDIAN);
br_add_callback(bitstream, ogg_crc, &checksum);
}
if (!setjmp(*br_try(bitstream))) {
do {
checksum = 0;
if (verifymodule_read_ogg_header(bitstream, &header) == OK) {
if (data_buffer_size < header.segment_length_total) {
data_buffer = realloc(data_buffer,
header.segment_length_total);
data_buffer_size = header.segment_length_total;
}
if (fread(data_buffer,
sizeof(uint8_t),
header.segment_length_total,
bitstream->input.file) !=
header.segment_length_total) {
PyErr_SetString(PyExc_IOError, "I/O error reading stream");
goto error;
}
for (i = 0; i < header.segment_length_total; i++)
ogg_crc(data_buffer[i], &checksum);
if (header.checksum != checksum) {
PyErr_SetString(PyExc_ValueError,
"checksum mismatch in stream");
goto error;
}
/* printf("calculated checksum : 0x%8.8X\n", checksum); */
if (has_previous_header) {
if (header.bitstream_serial_number !=
previous_header.bitstream_serial_number) {
PyErr_SetString(PyExc_ValueError,
"differing serial numbers in stream");
goto error;
}
if (header.page_sequence_number !=
(previous_header.page_sequence_number + 1)) {
PyErr_SetString(PyExc_ValueError,
"page sequence number not incrementing");
goto error;
}
previous_header = header;
} else {
previous_header = header;
has_previous_header = 1;
}
} else {
goto error;
}
} while (!(header.type & 0x4));
} else {
PyErr_SetString(PyExc_IOError, "I/O error reading stream");
goto error;
}
br_etry(bitstream);
free(data_buffer);
bitstream->input.file = NULL;
bitstream->free(bitstream);
Py_INCREF(Py_None);
return Py_None;
error:
br_etry(bitstream);
if (data_buffer != NULL)
free(data_buffer);
bitstream->input.file = NULL;
bitstream->free(bitstream);
return NULL;
}
static PyObject *
__read_dcd_header(PyObject *self, PyObject *args)
{
/* At this point we assume the file has been checked for existence and opened, and the DCD class
* contains a reference to the file (which can be used to retrieve the file descriptor)
*/
PyObject *temp;
fio_fd fd;
int rc;
char *remarks = NULL;
int len_remarks = 0;
dcdhandle *dcd = NULL;
if (! self) {
/* we were in fact called as a module function, try to retrieve
a matching object from args */
if( !PyArg_ParseTuple(args, "O", &self) )
return NULL;
} else {
/* we were obviously called as an object method so args should
only have the int value. */
if( !PyArg_ParseTuple(args, "") )
return NULL;
}
// Get the file object from the class
if (!PyObject_HasAttrString(self, "dcdfile")) {
// Raise exception
PyErr_SetString(PyExc_AttributeError, "dcdfile is not an attribute");
return NULL;
}
if ((temp = PyObject_GetAttrString(self, "dcdfile")) == NULL) { // This gives me a New Reference
// Raise exception
PyErr_SetString(PyExc_AttributeError, "dcdfile is not an attribute");
return NULL;
}
if (!PyFile_CheckExact(temp)) {
// Raise exception
PyErr_SetString(PyExc_TypeError, "dcdfile does not refer to a file object");
goto error;
}
fd = fileno(PyFile_AsFile(temp));
// No longer need the reference to temp
Py_DECREF(temp);
dcd = (dcdhandle *)malloc(sizeof(dcdhandle));
memset(dcd, 0, sizeof(dcdhandle));
dcd->fd = fd;
if ((rc = read_dcdheader(dcd->fd, &dcd->natoms, &dcd->nsets, &dcd->istart,
&dcd->nsavc, &dcd->delta, &dcd->nfixed, &dcd->freeind,
&dcd->fixedcoords, &dcd->reverse, &dcd->charmm, &remarks, &len_remarks))) {
// Raise exception
PyErr_SetString(PyExc_IOError, "Cannot read DCD header");
goto error;
}
/*
* Check that the file is big enough to really hold all the frames it claims to have
*/
{
off_t ndims, firstframesize, framesize, extrablocksize;
off_t filesize;
struct stat stbuf;
extrablocksize = dcd->charmm & DCD_HAS_EXTRA_BLOCK ? 48 + 8 : 0;
ndims = dcd->charmm & DCD_HAS_4DIMS ? 4 : 3;
firstframesize = (dcd->natoms+2) * ndims * sizeof(float) + extrablocksize;
framesize = (dcd->natoms-dcd->nfixed+2) * ndims * sizeof(float)
+ extrablocksize;
// Get filesize from file descriptor
memset(&stbuf, 0, sizeof(struct stat));
if (fstat(dcd->fd, &stbuf)) {
PyErr_SetString(PyExc_IOError, "Could not stat file");
goto error;
}
// Size of header
dcd->header_size = fio_ftell(dcd->fd);
/*
* It's safe to use ftell, even though ftell returns a long, because the
* header size is < 4GB.
*/
filesize = stbuf.st_size - fio_ftell(dcd->fd) - firstframesize;
if (filesize < 0) {
PyErr_SetString(PyExc_IOError, "DCD file appears to contain no timesteps");
goto error;
}
dcd->nsets = filesize / framesize + 1;
dcd->setsread = 0;
}
// We are at the first frame
dcd->first = 1;
temp = Py_BuildValue("s#", remarks, len_remarks);
if (temp == NULL) goto error;
if (PyObject_SetAttrString(self, "remarks", temp) == -1) {
PyErr_SetString(PyExc_AttributeError, "Could not create attribute remarks");
goto error;
}
Py_DECREF(temp);
temp = Py_BuildValue("i", dcd->natoms);
if (temp == NULL) goto error;
if (PyObject_SetAttrString(self, "n_atoms", temp) == -1) {
PyErr_SetString(PyExc_AttributeError, "Could not create attribute n_atoms");
goto error;
}
Py_DECREF(temp);
temp = Py_BuildValue("i", dcd->nsets);
if (temp == NULL) goto error;
if (PyObject_SetAttrString(self, "n_frames", temp) == -1) {
PyErr_SetString(PyExc_AttributeError, "Could not create attribute n_frames");
goto error;
}
Py_DECREF(temp);
temp = Py_BuildValue("i", dcd->nfixed);
if (temp == NULL) goto error;
if (PyObject_SetAttrString(self, "fixed", temp) == -1) {
PyErr_SetString(PyExc_AttributeError, "Could not create attribute fixed");
goto error;
}
Py_DECREF(temp);
temp = Py_BuildValue("i", dcd->istart);
if (temp == NULL) goto error;
if (PyObject_SetAttrString(self, "start_timestep", temp) == -1) {
PyErr_SetString(PyExc_AttributeError, "Could not create attribute fixed");
goto error;
}
Py_DECREF(temp);
temp = Py_BuildValue("i", dcd->nsavc);
if (temp == NULL) goto error;
if (PyObject_SetAttrString(self, "skip_timestep", temp) == -1) {
PyErr_SetString(PyExc_AttributeError, "Could not create attribute fixed");
goto error;
}
Py_DECREF(temp);
temp = Py_BuildValue("d", dcd->delta);
if (temp == NULL) goto error;
if (PyObject_SetAttrString(self, "delta", temp) == -1) {
PyErr_SetString(PyExc_AttributeError, "Could not create attribute fixed");
goto error;
}
Py_DECREF(temp);
temp = PyCObject_FromVoidPtr(dcd, NULL);
if (temp == NULL) goto error;
if (PyObject_SetAttrString(self, "_dcd_C_ptr", temp) == -1) {
PyErr_SetString(PyExc_AttributeError, "Could not create attribute _dcd_C_ptr");
goto error;
}
if ((dcd->charmm & DCD_IS_CHARMM) && (dcd->charmm & DCD_HAS_EXTRA_BLOCK)) {
dcd->with_unitcell = 1;
temp = Py_True;
} else {
temp = Py_False;
}
Py_INCREF(temp);
if (PyObject_SetAttrString(self, "periodic", temp) == -1) {
PyErr_SetString(PyExc_AttributeError, "Could not create attribute periodic");
goto error;
}
Py_DECREF(temp);
// For debugging purposes
temp = PyBuffer_FromMemory(dcd, sizeof(dcdhandle));
if (temp == NULL) goto error;
if (PyObject_SetAttrString(self, "_dcd_C_str", temp) == -1) {
PyErr_SetString(PyExc_AttributeError, "Could not create attribute _dcd_C_str");
goto error;
}
Py_DECREF(temp);
Py_INCREF(Py_None);
return Py_None;
error:
Py_XDECREF(temp);
if (dcd != NULL) free(dcd);
if (remarks != NULL) free(remarks);
return NULL;
}
