Save PNG in correct orientation with imsave

naoyak · naoyak · commit d9107888d3e3 · 2017-01-03T23:14:08.000-08:00
diff --git a/lib/matplotlib/image.py b/lib/matplotlib/image.py
@@ -521,7 +521,8 @@ def contains(self, mouseevent):
 
     def write_png(self, fname):
         """Write the image to png file with fname"""
-        im = self.to_rgba(self._A, bytes=True, norm=True)
+        im = self.to_rgba(self._A[::-1] if self.origin == 'lower' else self._A,
+                          bytes=True, norm=True)
         _png.write_png(im, fname)
 
     def set_data(self, A):
diff --git a/lib/matplotlib/tests/test_image.py b/lib/matplotlib/tests/test_image.py
@@ -150,21 +150,24 @@ def test_imsave_color_alpha():
     # acceptably preserved through a save/read roundtrip.
     from numpy import random
     random.seed(1)
-    data = random.rand(16, 16, 4)
 
-    buff = io.BytesIO()
-    plt.imsave(buff, data)
+    for origin in ['lower', 'upper']:
+        data = random.rand(16, 16, 4)
+        buff = io.BytesIO()
+        plt.imsave(buff, data, origin=origin)
 
-    buff.seek(0)
-    arr_buf = plt.imread(buff)
+        buff.seek(0)
+        arr_buf = plt.imread(buff)
 
-    # Recreate the float -> uint8 conversion of the data
-    # We can only expect to be the same with 8 bits of precision,
-    # since that's what the PNG file used.
-    data = (255*data).astype('uint8')
-    arr_buf = (255*arr_buf).astype('uint8')
+        # Recreate the float -> uint8 conversion of the data
+        # We can only expect to be the same with 8 bits of precision,
+        # since that's what the PNG file used.
+        data = (255*data).astype('uint8')
+        if origin == 'lower':
+            data = data[::-1]
+        arr_buf = (255*arr_buf).astype('uint8')
 
-    assert_array_equal(data, arr_buf)
+        assert_array_equal(data, arr_buf)
 
 @image_comparison(baseline_images=['image_alpha'], remove_text=True)
 def test_image_alpha():
