Add 16-bit PNG support to _png module.

mdboom · mdboom · commit 0e4f3ccbffe6 · 2012-06-05T14:21:36.000-04:00
diff --git a/lib/matplotlib/testing/compare.py b/lib/matplotlib/testing/compare.py
@@ -210,8 +210,8 @@ def compare_images( expected, actual, tol, in_decorator=False ):
       expected = convert(expected)
 
    # open the image files and remove the alpha channel (if it exists)
-   expectedImage = _png.read_png_uint8( expected )
-   actualImage = _png.read_png_uint8( actual )
+   expectedImage = _png.read_png_int( expected )
+   actualImage = _png.read_png_int( actual )
 
    actualImage, expectedImage = crop_to_same(actual, actualImage, expected, expectedImage)
 
diff --git a/lib/matplotlib/tests/baseline_images/test_png/pngtest16rgba.png b/lib/matplotlib/tests/baseline_images/test_png/pngtest16rgba.png
diff --git a/lib/matplotlib/tests/test_png.py b/lib/matplotlib/tests/test_png.py
@@ -3,6 +3,7 @@
 import matplotlib.cm as cm
 import glob
 import os
+import numpy as np
 
 @image_comparison(baseline_images=['pngsuite'], extensions=['png'])
 def test_pngsuite():
@@ -25,3 +26,12 @@ def test_pngsuite():
 
     plt.gca().get_frame().set_facecolor("#ddffff")
     plt.gca().set_xlim(0, len(files))
+
+
+def test_imread_png_uint16():
+    from matplotlib import _png
+    img = _png.read_png_int(os.path.join(os.path.dirname(__file__),
+                                         'baseline_images/test_png/pngtest16rgba.png'))
+    assert (img.dtype == np.uint16)
+
+    assert np.sum(img.flatten()) == 104855776
diff --git a/src/_png.cpp b/src/_png.cpp
@@ -48,6 +48,8 @@ class _png_module : public Py::ExtensionModule<_png_module>
                            "read_png_float(fileobj)");
         add_varargs_method("read_png_uint8", &_png_module::read_png_uint8,
                            "read_png_uint8(fileobj)");
+        add_varargs_method("read_png_int", &_png_module::read_png_int,
+                           "read_png_int(fileobj)");
         initialize("Module to write PNG files");
     }
 
@@ -57,7 +59,8 @@ class _png_module : public Py::ExtensionModule<_png_module>
     Py::Object write_png(const Py::Tuple& args);
     Py::Object read_png_uint8(const Py::Tuple& args);
     Py::Object read_png_float(const Py::Tuple& args);
-    PyObject* _read_png(const Py::Object& py_fileobj, const bool float_result);
+    Py::Object read_png_int(const Py::Tuple& args);
+    PyObject* _read_png(const Py::Object& py_fileobj, const bool float_result, const int bit_depth = -1);
 };
 
 static void write_png_data(png_structp png_ptr, png_bytep data, png_size_t length)
@@ -297,7 +300,8 @@ static void read_png_data(png_structp png_ptr, png_bytep data, png_size_t length
 }
 
 PyObject*
-_png_module::_read_png(const Py::Object& py_fileobj, const bool float_result)
+_png_module::_read_png(const Py::Object& py_fileobj, const bool float_result,
+                       const int result_bit_depth)
 {
     png_byte header[8];   // 8 is the maximum size that can be checked
     FILE* fp = NULL;
@@ -502,7 +506,18 @@ _png_module::_read_png(const Py::Object& py_fileobj, const bool float_result)
             }
         }
     } else {
-        A = (PyArrayObject *) PyArray_SimpleNew(num_dims, dimensions, NPY_UBYTE);
+        if (result_bit_depth == 8) {
+            A = (PyArrayObject *) PyArray_SimpleNew(num_dims, dimensions, NPY_UBYTE);
+        } else {
+            if (bit_depth == 8) {
+                A = (PyArrayObject *) PyArray_SimpleNew(num_dims, dimensions, NPY_UBYTE);
+            } else if (bit_depth == 16) {
+                A = (PyArrayObject *) PyArray_SimpleNew(num_dims, dimensions, NPY_UINT16);
+            } else {
+                throw Py::RuntimeError(
+                    "_image_module::readpng: image has unknown bit depth");
+            }
+        }
 
         if (A == NULL)
         {
@@ -518,9 +533,17 @@ _png_module::_read_png(const Py::Object& py_fileobj, const bool float_result)
                 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++)
-                    {
-                        *(png_byte*)(A->data + offset + p*A->strides[2]) = ptr[p] >> 8;
+
+                    if (bit_depth == 16) {
+                        for (png_uint_32 p = 0; p < (png_uint_32)dimensions[2]; p++)
+                        {
+                            *(png_uint_16*)(A->data + offset + p*A->strides[2]) = ptr[p];
+                        }
+                    } else {
+                        for (png_uint_32 p = 0; p < (png_uint_32)dimensions[2]; p++)
+                        {
+                            *(png_byte*)(A->data + offset + p*A->strides[2]) = ptr[p] >> 8;
+                        }
                     }
                 }
                 else
@@ -569,6 +592,12 @@ _png_module::read_png_float(const Py::Tuple& args)
 
 Py::Object
 _png_module::read_png_uint8(const Py::Tuple& args)
+{
+    throw Py::RuntimeError("read_png_uint8 is deprecated.  Use read_png_int instead.");
+}
+
+Py::Object
+_png_module::read_png_int(const Py::Tuple& args)
 {
     args.verify_length(1);
     return Py::asObject(_read_png(args[0], false));
