add webp support to px.imshow

Signed-off-by: maximsmol <[email protected]>

Author: maximsmol

packages/python/plotly/_plotly_utils/data_utils.py

@@ -10,7 +10,9 @@
     pil_imported = False
 
 
-def image_array_to_data_uri(img, backend="pil", compression=4, ext="png"):
+def image_array_to_data_uri(
+    img, backend="pil", compression=4, ext="webp", backend_kwargs=None
+):
     """Converts a numpy array of uint8 into a base64 png or jpg string.
 
     Parameters
@@ -22,8 +24,10 @@ def image_array_to_data_uri(img, backend="pil", compression=4, ext="png"):
         otherwise pypng.
     compression: int, between 0 and 9
         compression level to be passed to the backend
-    ext: str, 'png' or 'jpg'
+    ext: str, 'webp', 'png', or 'jpg'
         compression format used to generate b64 string
+    backend_kwargs : dict or None
+        keyword arguments to be passed to the backend
     """
     # PIL and pypng error messages are quite obscure so we catch invalid compression values
     if compression < 0 or compression > 9:
@@ -41,15 +45,25 @@ def image_array_to_data_uri(img, backend="pil", compression=4, ext="png"):
     if backend == "auto":
         backend = "pil" if pil_imported else "pypng"
     if ext != "png" and backend != "pil":
-        raise ValueError("jpg binary strings are only available with PIL backend")
+        raise ValueError(
+            "webp and jpg binary strings are only available with PIL backend"
+        )
+
+    if backend_kwargs is None:
+        backend_kwargs = {}
 
     if backend == "pypng":
         ndim = img.ndim
         sh = img.shape
         if ndim == 3:
             img = img.reshape((sh[0], sh[1] * sh[2]))
         w = Writer(
-            sh[1], sh[0], greyscale=(ndim == 2), alpha=alpha, compression=compression
+            sh[1],
+            sh[0],
+            greyscale=(ndim == 2),
+            alpha=alpha,
+            compression=compression,
+            **backend_kwargs
         )
         img_png = from_array(img, mode=mode)
         prefix = "data:image/png;base64,"
@@ -63,13 +77,22 @@ def image_array_to_data_uri(img, backend="pil", compression=4, ext="png"):
                 "install pillow or use `backend='pypng'."
             )
         pil_img = Image.fromarray(img)
-        if ext == "jpg" or ext == "jpeg":
+        if ext == "webp":
+            prefix = "data:image/webp;base64,"
+            ext = "webp"
+        elif ext == "jpg" or ext == "jpeg":
             prefix = "data:image/jpeg;base64,"
             ext = "jpeg"
         else:
             prefix = "data:image/png;base64,"
             ext = "png"
         with BytesIO() as stream:
-            pil_img.save(stream, format=ext, compress_level=compression)
+            pil_img.save(
+                stream,
+                format=ext,
+                compress_level=compression,
+                lossless=True,
+                **backend_kwargs
+            )
             base64_string = prefix + base64.b64encode(stream.getvalue()).decode("utf-8")
     return base64_string

packages/python/plotly/plotly/express/_imshow.py

@@ -78,7 +78,8 @@ def imshow(
     binary_string=None,
     binary_backend="auto",
     binary_compression_level=4,
-    binary_format="png",
+    binary_format="webp",
+    binary_backend_kwargs=None,
     text_auto=False,
 ) -> go.Figure:
     """
@@ -204,10 +205,15 @@ def imshow(
         test `len(fig.data[0].source)` and to time the execution of `imshow` to
         tune the level of compression. 0 means no compression (not recommended).
 
-    binary_format: str, 'png' (default) or 'jpg'
-        compression format used to generate b64 string. 'png' is recommended
-        since it uses lossless compression, but 'jpg' (lossy) compression can
-        result if smaller binary strings for natural images.
+    binary_format: str, 'webp' (default), 'png', or 'jpg'
+        compression format used to generate b64 string. 'webp' is recommended
+        since it supports both lossless and lossy compression with better quality
+        then 'png' or 'jpg' of similar sizes, but 'jpg' or 'png' can be used for
+        environments that do not support 'webp'.
+
+    binary_backend_kwargs : dict or None
+        keyword arguments for the image backend. For Pillow, these are passed to `Image.save`.
+        For 'pypng', these are passed to `Writer.__init__`
 
     text_auto: bool or str (default `False`)
         If `True` or a string, single-channel `img` values will be displayed as text.
@@ -502,6 +508,7 @@ def imshow(
                     backend=binary_backend,
                     compression=binary_compression_level,
                     ext=binary_format,
+                    binary_backend_kwargs=binary_backend_kwargs,
                 )
                 for index_tup in itertools.product(*iterables)
             ]