Merge pull request #1460 from jni/fix-pil-imsave

Allow keyword arguments to imsave for PIL plugin

skimage/_shared/_tempfile.py

@@ -0,0 +1,27 @@
+from tempfile import NamedTemporaryFile
+from contextlib import contextmanager
+import os
+
+@contextmanager
+def temporary_file(suffix=''):
+    """Yield a writeable temporary filename that is deleted on context exit.
+
+    Parameters
+    ----------
+    suffix : string, optional
+        The suffix for the file.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> from skimage import io
+    >>> with temporary_file('.tif') as tempfile:
+    ...     im = np.zeros((5, 5), np.uint8)
+    ...     io.imsave(tempfile, im)
+    ...     assert np.all(io.imread(tempfile) == im)
+    """
+    tempfile_stream = NamedTemporaryFile(suffix=suffix, delete=False)
+    tempfile = tempfile_stream.name
+    tempfile_stream.close()
+    yield tempfile
+    os.remove(tempfile)

skimage/io/_plugins/pil_plugin.py

@@ -196,7 +196,7 @@ def ndarray_to_pil(arr, format_str=None):
     return im
 
 
-def imsave(fname, arr, format_str=None):
+def imsave(fname, arr, format_str=None, **kwargs):
     """Save an image to disk.
 
     Parameters
@@ -210,6 +210,12 @@ def imsave(fname, arr, format_str=None):
     format_str: str
         Format to save as, this is defaulted to PNG if using a file-like
         object; this will be derived from the extension if fname is a string
+    kwargs: dict
+        Keyword arguments to the Pillow save function (or tifffile save
+        function, for Tiff files). These are format dependent. For example,
+        Pillow's JPEG save function supports an integer ``quality`` argument
+        with values in [1, 95], while TIFFFile supports a ``compress``
+        integer argument with values in [0, 9].
 
     Notes
     -----
@@ -251,7 +257,7 @@ def imsave(fname, arr, format_str=None):
             use_tif = True
 
     if use_tif:
-        tif_imsave(fname, arr)
+        tif_imsave(fname, arr, **kwargs)
         return
 
     if arr.ndim not in (2, 3):
@@ -262,4 +268,4 @@ def imsave(fname, arr, format_str=None):
             raise ValueError("Invalid number of channels in image array.")
 
     img = ndarray_to_pil(arr, format_str=format_str)
-    img.save(fname, format=format_str)
+    img.save(fname, format=format_str, **kwargs)

skimage/io/tests/test_pil.py

@@ -11,6 +11,7 @@ from .. import (imread, imsave, use_plugin, reset_plugins,
                         Image as ioImage)
 from ..._shared.testing import mono_check, color_check
 from ..._shared._warnings import expected_warnings
+from ..._shared._tempfile import temporary_file
 
 from six import BytesIO
 
@@ -97,6 +98,16 @@ def test_imread_truncated_jpg():
                   os.path.join(data_dir, 'truncated.jpg'))
 
 
+def test_jpg_quality_arg():
+    chessboard = np.load(os.path.join(data_dir, 'chessboard_GRAY_U8.npy'))
+    with temporary_file(suffix='.jpg') as jpg:
+        imsave(jpg, chessboard, quality=95)
+        im = imread(jpg)
+        sim = ssim(chessboard, im,
+                   dynamic_range=chessboard.max() - chessboard.min())
+        assert sim > 0.99
+
+
 def test_imread_uint16_big_endian():
     expected = np.load(os.path.join(data_dir, 'chessboard_GRAY_U8.npy'))
     img = imread(os.path.join(data_dir, 'chessboard_GRAY_U16B.tif'))
@@ -106,12 +117,10 @@ def test_imread_uint16_big_endian():
 
 class TestSave:
     def roundtrip_file(self, x):
-        f = NamedTemporaryFile(suffix='.png')
-        fname = f.name
-        f.close()
-        imsave(fname, x)
-        y = imread(fname)
-        return y
+        with temporary_file(suffix='.png') as fname:
+            imsave(fname, x)
+            y = imread(fname)
+            return y
 
     def roundtrip_pil_image(self, x):
         pil_image = ndarray_to_pil(x)
@@ -214,25 +223,33 @@ def test_cmyk():
 
 
 class TestSaveTIF:
-    def roundtrip(self, dtype, x):
-        f = NamedTemporaryFile(suffix='.tif')
-        fname = f.name
-        f.close()
-        imsave(fname, x)
-        y = imread(fname)
-        assert_array_equal(x, y)
+    def roundtrip(self, dtype, x, compress):
+        with temporary_file(suffix='.tif') as fname:
+            if dtype == np.bool:
+                expected = ['low contrast']
+            else:
+                expected = []
+            with expected_warnings(expected):
+                if compress > 0:
+                    imsave(fname, x, compress=compress)
+                else:
+                    imsave(fname, x)
+            y = imread(fname)
+            assert_array_equal(x, y)
 
     def test_imsave_roundtrip(self):
         for shape in [(10, 10), (10, 10, 3), (10, 10, 4)]:
             for dtype in (np.uint8, np.uint16, np.int16, np.float32,
                           np.float64, np.bool):
-                x = np.random.rand(*shape)
+                for compress in [0, 2]:
+                    x = np.random.rand(*shape)
 
-                if not np.issubdtype(dtype, float) and not dtype == np.bool:
-                    x = (x * np.iinfo(dtype).max).astype(dtype)
-                else:
-                    x = x.astype(dtype)
-                yield self.roundtrip, dtype, x
+                    if not np.issubdtype(dtype, float) and \
+                                                        not dtype == np.bool:
+                        x = (x * np.iinfo(dtype).max).astype(dtype)
+                    else:
+                        x = x.astype(dtype)
+                    yield self.roundtrip, dtype, x, compress
 
 if __name__ == "__main__":
     run_module_suite()