Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 7bde3b8

Browse files
mdboommdboom
committed
Support Python file-like objects in readpng
svn path=/trunk/matplotlib/; revision=7909
1 parent a0c08db commit 7bde3b8

3 files changed

Lines changed: 131 additions & 46 deletions

File tree

‎lib/matplotlib/image.py‎

Lines changed: 23 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -938,9 +938,14 @@ def draw(self, renderer, *args, **kwargs):
938938

939939

940940

941-
def imread(fname):
941+
def imread(fname, format=None):
942942
"""
943-
Return image file in *fname* as :class:`numpy.array`.
943+
Return image file in *fname* as :class:`numpy.array`. *fname* may
944+
be a string path or a Python file-like object.
945+
946+
If *format* is provided, will try to read file of that type,
947+
otherwise the format is deduced from the filename. If nothing can
948+
be deduced, PNG is tried.
944949
945950
Return value is a :class:`numpy.array`. For grayscale images, the
946951
return array is MxN. For RGB images, the return value is MxNx3.
@@ -959,11 +964,15 @@ def pilread():
959964
image = Image.open( fname )
960965
return pil_to_array(image)
961966

962-
963-
handlers = {'png' :_png.read_png,
964-
}
965-
basename, ext = os.path.splitext(fname)
966-
ext = ext.lower()[1:]
967+
handlers = {'png' :_png.read_png, }
968+
if format is None:
969+
if cbook.is_string_like(fname):
970+
basename, ext = os.path.splitext(fname)
971+
ext = ext.lower()[1:]
972+
else:
973+
ext = 'png'
974+
else:
975+
ext = format
967976

968977
if ext not in handlers.keys():
969978
im = pilread()
@@ -972,6 +981,13 @@ def pilread():
972981
return im
973982

974983
handler = handlers[ext]
984+
985+
# To handle Unicode filenames, we pass a file object to the PNG
986+
# reader extension, since Python handles them quite well, but it's
987+
# tricky in C.
988+
if cbook.is_string_like(fname):
989+
fname = open(fname, 'rb')
990+
975991
return handler(fname)
976992

977993

‎lib/matplotlib/tests/test_image.py‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,9 @@
44
import matplotlib.pyplot as plt
55
from nose.tools import assert_raises
66

7+
import cStringIO
8+
import os
9+
710
@image_comparison(baseline_images=['image_interps'])
811
def test_image_interps():
912
'make the basic nearest, bilinear and bicubic interps'
@@ -26,6 +29,24 @@ def test_image_interps():
2629

2730
fig.savefig('image_interps')
2831

32+
def test_image_python_io():
33+
fig = plt.figure()
34+
ax = fig.add_subplot(111)
35+
ax.plot([1,2,3])
36+
buffer = cStringIO.StringIO()
37+
fig.savefig(buffer)
38+
buffer.seek(0)
39+
plt.imread(buffer)
40+
41+
def test_image_unicode_io():
42+
fig = plt.figure()
43+
ax = fig.add_subplot(111)
44+
ax.plot([1,2,3])
45+
fname = u"\u0a3a\u0a3a.png"
46+
fig.savefig(fname)
47+
plt.imread(fname)
48+
os.remove(fname)
49+
2950
if __name__=='__main__':
3051
import nose
3152
nose.runmodule(argv=['-s','--with-doctest'], exit=False)

‎src/_png.cpp‎

Lines changed: 87 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -130,12 +130,12 @@ Py::Object _png_module::write_png(const Py::Tuple& args)
130130
png_init_io(png_ptr, fp);
131131
} else {
132132
png_set_write_fn(png_ptr, (void*)py_fileobj.ptr(),
133-
&write_png_data, &flush_png_data);
133+
&write_png_data, &flush_png_data);
134134
}
135135
png_set_IHDR(png_ptr, info_ptr,
136-
width, height, 8,
137-
PNG_COLOR_TYPE_RGB_ALPHA, PNG_INTERLACE_NONE,
138-
PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);
136+
width, height, 8,
137+
PNG_COLOR_TYPE_RGB_ALPHA, PNG_INTERLACE_NONE,
138+
PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);
139139

140140
// Save the dpi of the image in the file
141141
if (args.size() == 5) {
@@ -157,14 +157,14 @@ Py::Object _png_module::write_png(const Py::Tuple& args)
157157
png_write_image(png_ptr, row_pointers);
158158
png_write_end(png_ptr, info_ptr);
159159
} catch (...) {
160-
if (fp && close_file) fclose(fp);
161-
delete [] row_pointers;
162-
/* Changed calls to png_destroy_write_struct to follow
163-
http://www.libpng.org/pub/png/libpng-manual.txt.
164-
This ensures the info_ptr memory is released.
165-
*/
166-
if (png_ptr && info_ptr) png_destroy_write_struct(&png_ptr, &info_ptr);
167-
throw;
160+
if (fp && close_file) fclose(fp);
161+
delete [] row_pointers;
162+
/* Changed calls to png_destroy_write_struct to follow
163+
http://www.libpng.org/pub/png/libpng-manual.txt.
164+
This ensures the info_ptr memory is released.
165+
*/
166+
if (png_ptr && info_ptr) png_destroy_write_struct(&png_ptr, &info_ptr);
167+
throw;
168168
}
169169

170170
png_destroy_write_struct(&png_ptr, &info_ptr);
@@ -174,39 +174,85 @@ Py::Object _png_module::write_png(const Py::Tuple& args)
174174
return Py::Object();
175175
}
176176

177+
static void _read_png_data(PyObject* py_file_obj, png_bytep data, png_size_t length) {
178+
PyObject* read_method = PyObject_GetAttrString(py_file_obj, "read");
179+
PyObject* result = NULL;
180+
char *buffer;
181+
Py_ssize_t bufflen;
182+
if (read_method)
183+
result = PyObject_CallFunction(read_method, (char *)"i", length);
184+
if (PyString_AsStringAndSize(result, &buffer, &bufflen) == 0) {
185+
if (bufflen == (Py_ssize_t)length) {
186+
memcpy(data, buffer, length);
187+
}
188+
}
189+
Py_XDECREF(read_method);
190+
Py_XDECREF(result);
191+
}
192+
193+
static void read_png_data(png_structp png_ptr, png_bytep data, png_size_t length) {
194+
PyObject* py_file_obj = (PyObject*)png_get_io_ptr(png_ptr);
195+
_read_png_data(py_file_obj, data, length);
196+
}
177197

178198
Py::Object
179199
_png_module::read_png(const Py::Tuple& args) {
180200

181201
args.verify_length(1);
182-
std::string fname = Py::String(args[0]);
183-
184-
png_byte header[8]; // 8 is the maximum size that can be checked
202+
png_byte header[8]; // 8 is the maximum size that can be checked
203+
FILE* fp = NULL;
204+
bool close_file = false;
185205

186-
FILE *fp = fopen(fname.c_str(), "rb");
187-
if (!fp)
188-
throw Py::RuntimeError(Printf("_image_module::readpng could not open PNG file %s for reading", fname.c_str()).str());
206+
Py::Object py_fileobj = Py::Object(args[0]);
207+
if (py_fileobj.isString()) {
208+
std::string fileName = Py::String(py_fileobj);
209+
const char *file_name = fileName.c_str();
210+
if ((fp = fopen(file_name, "rb")) == NULL)
211+
throw Py::RuntimeError( Printf("Could not open file %s for reading", file_name).str() );
212+
close_file = true;
213+
} else if (PyFile_CheckExact(py_fileobj.ptr())) {
214+
fp = PyFile_AsFile(py_fileobj.ptr());
215+
} else {
216+
PyObject* read_method = PyObject_GetAttrString(py_fileobj.ptr(), "read");
217+
if (!(read_method && PyCallable_Check(read_method))) {
218+
Py_XDECREF(read_method);
219+
throw Py::TypeError("Object does not appear to be a 8-bit string path or a Python file-like object");
220+
}
221+
Py_XDECREF(read_method);
222+
}
189223

190-
if (fread(header, 1, 8, fp) != 8)
191-
throw Py::RuntimeError("_image_module::readpng: error reading PNG header");
192-
if (png_sig_cmp(header, 0, 8))
224+
if (fp) {
225+
if (fread(header, 1, 8, fp) != 8) {
226+
throw Py::RuntimeError("_image_module::readpng: error reading PNG header");
227+
}
228+
} else {
229+
_read_png_data(py_fileobj.ptr(), header, 8);
230+
}
231+
if (png_sig_cmp(header, 0, 8)) {
193232
throw Py::RuntimeError("_image_module::readpng: file not recognized as a PNG file");
194-
233+
}
195234

196235
/* initialize stuff */
197236
png_structp png_ptr = png_create_read_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);
198237

199-
if (!png_ptr)
238+
if (!png_ptr) {
200239
throw Py::RuntimeError("_image_module::readpng: png_create_read_struct failed");
240+
}
201241

202242
png_infop info_ptr = png_create_info_struct(png_ptr);
203-
if (!info_ptr)
243+
if (!info_ptr) {
204244
throw Py::RuntimeError("_image_module::readpng: png_create_info_struct failed");
245+
}
205246

206-
if (setjmp(png_jmpbuf(png_ptr)))
247+
if (setjmp(png_jmpbuf(png_ptr))) {
207248
throw Py::RuntimeError("_image_module::readpng: error during init_io");
249+
}
208250

209-
png_init_io(png_ptr, fp);
251+
if (fp) {
252+
png_init_io(png_ptr, fp);
253+
} else {
254+
png_set_read_fn(png_ptr, (void*)py_fileobj.ptr(), &read_png_data);
255+
}
210256
png_set_sig_bytes(png_ptr, 8);
211257
png_read_info(png_ptr, info_ptr);
212258

@@ -272,26 +318,28 @@ _png_module::read_png(const Py::Tuple& args) {
272318

273319
for (png_uint_32 y = 0; y < height; y++) {
274320
png_byte* row = row_pointers[y];
275-
for (png_uint_32 x = 0; x < width; x++) {
276-
size_t offset = y*A->strides[0] + x*A->strides[1];
277-
if (bit_depth == 16) {
278-
png_uint_16* ptr = &reinterpret_cast<png_uint_16*> (row)[x * dimensions[2]];
321+
for (png_uint_32 x = 0; x < width; x++) {
322+
size_t offset = y*A->strides[0] + x*A->strides[1];
323+
if (bit_depth == 16) {
324+
png_uint_16* ptr = &reinterpret_cast<png_uint_16*> (row)[x * dimensions[2]];
325+
for (png_uint_32 p = 0; p < (png_uint_32)dimensions[2]; p++)
326+
*(float*)(A->data + offset + p*A->strides[2]) = (float)(ptr[p]) / max_value;
327+
} else {
328+
png_byte* ptr = &(row[x * dimensions[2]]);
279329
for (png_uint_32 p = 0; p < (png_uint_32)dimensions[2]; p++)
280-
*(float*)(A->data + offset + p*A->strides[2]) = (float)(ptr[p]) / max_value;
281-
} else {
282-
png_byte* ptr = &(row[x * dimensions[2]]);
283-
for (png_uint_32 p = 0; p < (png_uint_32)dimensions[2]; p++)
284-
{
285-
*(float*)(A->data + offset + p*A->strides[2]) = (float)(ptr[p]) / max_value;
286-
}
287-
}
330+
{
331+
*(float*)(A->data + offset + p*A->strides[2]) = (float)(ptr[p]) / max_value;
332+
}
333+
}
288334
}
289335
}
290336

291337
//free the png memory
292338
png_read_end(png_ptr, info_ptr);
293339
png_destroy_read_struct(&png_ptr, &info_ptr, png_infopp_NULL);
294-
fclose(fp);
340+
if (close_file) {
341+
fclose(fp);
342+
}
295343
for (row = 0; row < height; row++)
296344
delete [] row_pointers[row];
297345
delete [] row_pointers;

0 commit comments

Comments
 (0)