add jpg compression option

Author: emmanuelle

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

@@ -24,7 +24,7 @@
 _float_types = []
 
 
-def _array_to_b64str(img, backend="pil", compression=4):
+def _array_to_b64str(img, backend="pil", compression=4, ext="png"):
     """Converts a numpy array of uint8 into a base64 png string.
 
     Parameters
@@ -36,6 +36,8 @@ def _array_to_b64str(img, backend="pil", compression=4):
         otherwise pypng.
     compression: int, between 0 and 9
         compression level to be passed to the backend
+    ext: str, 'png' or 'jpg'
+        compression format used to generate b64 string
     """
     # PIL and pypng error messages are quite obscure so we catch invalid compression values
     if compression < 0 or compression > 9:
@@ -52,6 +54,9 @@ def _array_to_b64str(img, backend="pil", compression=4):
         raise ValueError("Invalid image shape")
     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")
+
     if backend == "pypng":
         ndim = img.ndim
         sh = img.shape
@@ -72,9 +77,14 @@ def _array_to_b64str(img, backend="pil", compression=4):
                 "install pillow or use `backend='pypng'."
             )
         pil_img = Image.fromarray(img)
-        prefix = "data:image/png;base64,"
+        if 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="png", compress_level=compression)
+            pil_img.save(stream, format=ext, compress_level=compression)
             base64_string = prefix + base64.b64encode(stream.getvalue()).decode("utf-8")
     return base64_string
 
@@ -135,6 +145,7 @@ def imshow(
     binary_string=None,
     binary_backend="auto",
     binary_compression_level=4,
+    binary_format="png",
 ):
     """
     Display an image, i.e. data on a 2D regular raster.
@@ -239,6 +250,9 @@ 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
+
     Returns
     -------
     fig : graph_objects.Figure containing the displayed image
@@ -410,6 +424,7 @@ def imshow(
                 img_rescaled,
                 backend=binary_backend,
                 compression=binary_compression_level,
+                ext=binary_format,
             )
             trace = go.Image(source=img_str)
         else:

packages/python/plotly/plotly/tests/test_core/test_px/test_imshow.py

@@ -18,7 +18,7 @@ def decode_image_string(image_string):
     if "png" in image_string[:22]:
         return np.asarray(Image.open(BytesIO(base64.b64decode(image_string[22:]))))
     elif "jpeg" in image_string[:23]:
-        return np.asaray(Image.open(BytesIO(base64.b64decode(image_string[23:]))))
+        return np.asarray(Image.open(BytesIO(base64.b64decode(image_string[23:]))))
     else:
         raise ValueError("image string format not recognized")
 
@@ -61,12 +61,24 @@ def test_automatic_zmax_from_dtype():
 
 
 @pytest.mark.parametrize("binary_string", [False, True])
-def test_origin(binary_string):
+@pytest.mark.parametrize("binary_format", ["png", "jpg"])
+def test_origin(binary_string, binary_format):
     for i, img in enumerate([img_rgb, img_gray]):
-        fig = px.imshow(img, origin="lower", binary_string=binary_string)
+        fig = px.imshow(
+            img,
+            origin="lower",
+            binary_string=binary_string,
+            binary_format=binary_format,
+        )
         assert fig.layout.yaxis.autorange is True
-        if binary_string and i == 0:
+        if binary_string and i == 0 and binary_format == "png":
+            # The equality below does not hold for jpeg compression since it's lossy
             assert np.all(img[::-1] == decode_image_string(fig.data[0].source))
+        if binary_string:
+            if binary_format == "jpg":
+                assert fig.data[0].source[:15] == "data:image/jpeg"
+            else:
+                assert fig.data[0].source[:14] == "data:image/png"
     fig = px.imshow(img_rgb, binary_string=binary_string)
     assert fig.layout.yaxis.autorange is None
     fig = px.imshow(img_gray, binary_string=binary_string)