more tests and docstring

Author: emmanuelle

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

@@ -33,18 +33,20 @@
 _integer_ranges = {t: (np.iinfo(t).min, np.iinfo(t).max) for t in _integer_types}
 
 
-def _array_to_b64str(img, ext='png'):
+def _array_to_b64str(img, ext="png"):
     pil_img = Image.fromarray(img)
     if ext == "jpg":
-       ext = "jpeg"
+        ext = "jpeg"
     buff = BytesIO()
     pil_img.save(buff, format=ext)
-    if ext == 'png':
-        prefix = b'data:image/png;base64,'
-    elif ext == 'jpeg':
-        prefix = b'data:image/jpeg;base64,'
+    if ext == "png":
+        prefix = b"data:image/png;base64,"
+    elif ext == "jpeg":
+        prefix = b"data:image/jpeg;base64,"
     else:
-        raise ValueError("accepted image formats are 'png' and 'jpeg' but %s was passed" %format)
+        raise ValueError(
+            "accepted image formats are 'png' and 'jpeg' but %s was passed" % format
+        )
     image_string = (prefix + base64.b64encode(buff.getvalue())).decode("utf-8")
     return image_string
 
@@ -177,6 +179,22 @@ def imshow(
       - if None, 'equal' is used for numpy arrays and 'auto' for xarrays
         (which have typically heterogeneous coordinates)
 
+    use_binary_string: bool, default None
+        if True, the image data are first rescaled and encoded as uint8 and
+        then passed to plotly.js as a b64 PNG string. If False, data are passed
+        unchanged as a numerical array. Setting to True may lead to performance
+        gains, at the cost of a loss of precision depending on the original data
+        type. If None, use_binary_string is set to True for multichannel (eg) RGB
+        arrays, and to False for single-channel (2D) arrays. 2D arrays are
+        represented as grayscale and with no colorbar if use_binary_string is
+        True.
+
+    contrast_rescaling: 'image', 'dtype', or None
+        how to determine data values corresponding to the bounds of the color
+        range, when zmin or zmax are not passed. If `image`, the min and max
+        values of the image are used. If `dtype`, a heuristic based on the image
+        data type is used.
+
     Returns
     -------
     fig : graph_objects.Figure containing the displayed image
@@ -203,10 +221,10 @@ def imshow(
     if xarray_imported and isinstance(img, xarray.DataArray):
         if use_binary_string:
             raise ValueError(
-                    "It is not possible to use binary image strings for xarrays."
-                    "Please pass your data as a numpy array instead using"
-                    "`img.values`"
-                    )
+                "It is not possible to use binary image strings for xarrays."
+                "Please pass your data as a numpy array instead using"
+                "`img.values`"
+            )
         y_label, x_label = img.dims[0], img.dims[1]
         # np.datetime64 is not handled correctly by go.Heatmap
         for ax in [x_label, y_label]:
@@ -254,7 +272,9 @@ def imshow(
     else:
         has_nans = np.any(np.isnan(img))
         if has_nans and use_binary_string:
-            raise ValueError("Binary strings cannot be used with arrays containing NaNs")
+            raise ValueError(
+                "Binary strings cannot be used with arrays containing NaNs"
+            )
 
     # --------------- Starting from here img is always a numpy array --------
     img = np.asanyarray(img)
@@ -268,8 +288,8 @@ def imshow(
         img = 255 * img.astype(np.uint8)
 
     if contrast_rescaling is None:
-        contrast_rescaling='image' if img.ndim == 2 else 'dtype'
-    if contrast_rescaling == 'image':
+        contrast_rescaling = "image" if img.ndim == 2 else "dtype"
+    if contrast_rescaling == "image":
         if (zmin is not None or use_binary_string) and zmax is None:
             zmax = img.max()
         if (zmax is not None or use_binary_string) and zmin is None:
@@ -312,14 +332,30 @@ def imshow(
             layout["coloraxis1"]["colorbar"] = dict(title_text=labels["color"])
 
     # For 2D+RGB data, use Image trace
-    elif img.ndim == 3 and img.shape[-1] in [3, 4] or (img.ndim == 2 and use_binary_string):
+    elif (
+        img.ndim == 3
+        and img.shape[-1] in [3, 4]
+        or (img.ndim == 2 and use_binary_string)
+    ):
         zmin, zmax = _vectorize_zvalue(zmin), _vectorize_zvalue(zmax)
         if use_binary_string:
             if img.ndim == 2:
-                img_rescaled = rescale_intensity(img, in_range=(zmin[0], zmax[0]), out_range=np.uint8)
+                img_rescaled = rescale_intensity(
+                    img, in_range=(zmin[0], zmax[0]), out_range=np.uint8
+                )
             else:
-                img_rescaled = np.dstack([rescale_intensity(img[..., ch], in_range=(zmin[ch], zmax[ch]), out_range=np.uint8)
-                                            for ch in range(img.shape[-1])])
+                img_rescaled = np.dstack(
+                    [
+                        rescale_intensity(
+                            img[..., ch],
+                            in_range=(zmin[ch], zmax[ch]),
+                            out_range=np.uint8,
+                        )
+                        for ch in range(img.shape[-1])
+                    ]
+                )
+            if origin == "lower":
+                img_rescaled = img_rescaled[::-1]
             img_str = _array_to_b64str(img_rescaled)
             trace = go.Image(source=img_str)
         else:
@@ -343,11 +379,17 @@ def imshow(
         layout_patch["margin"] = {"t": 60}
     fig = go.Figure(data=trace, layout=layout)
     fig.update_layout(layout_patch)
-    if not use_binary_string:
-        fig.update_traces(
-            hovertemplate="%s: %%{x}<br>%s: %%{y}<br>%s: %%{z}<extra></extra>"
-            % (labels["x"] or "x", labels["y"] or "y", labels["color"] or "color",)
+    # does not work yet, Antoine working on it
+    hover_name = "z" if not use_binary_string else "colorLabel"
+    fig.update_traces(
+        hovertemplate="%s: %%{x}<br>%s: %%{y}<br>%s: %%{%s}<extra></extra>"
+        % (
+            labels["x"] or "x",
+            labels["y"] or "y",
+            labels["color"] or "color",
+            hover_name,
         )
+    )
     if labels["x"]:
         fig.update_xaxes(title_text=labels["x"])
     if labels["y"]:

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

@@ -15,9 +15,9 @@ def decode_image_string(image_string):
     """
     Converts image string to numpy array.
     """
-    if 'png' in image_string[:22]:
+    if "png" in image_string[:22]:
         return np.asarray(Image.open(BytesIO(base64.b64decode(image_string[22:]))))
-    elif 'jpeg' in image_string[:23]:
+    elif "jpeg" in image_string[:23]:
         return np.asaray(Image.open(BytesIO(base64.b64decode(image_string[23:]))))
     else:
         raise ValueError("image string format not recognized")
@@ -57,14 +57,20 @@ def test_automatic_zmax_from_dtype():
         assert fig.data[0]["zmax"] == (val, val, val, 1)
 
 
-def test_origin():
-    for img in [img_rgb, img_gray]:
-        fig = px.imshow(img, origin="lower")
+@pytest.mark.parametrize("use_binary_string", [False, True])
+def test_origin(use_binary_string):
+    for i, img in enumerate([img_rgb, img_gray]):
+        fig = px.imshow(img, origin="lower", use_binary_string=use_binary_string)
         assert fig.layout.yaxis.autorange == True
-    fig = px.imshow(img_rgb)
+        if use_binary_string and i == 0:
+            assert np.all(img[::-1] == decode_image_string(fig.data[0].source))
+    fig = px.imshow(img_rgb, use_binary_string=use_binary_string)
     assert fig.layout.yaxis.autorange is None
-    fig = px.imshow(img_gray)
-    assert fig.layout.yaxis.autorange == "reversed"
+    fig = px.imshow(img_gray, use_binary_string=use_binary_string)
+    if use_binary_string:
+        assert fig.layout.yaxis.autorange is None
+    else:
+        assert fig.layout.yaxis.autorange == "reversed"
 
 
 def test_colorscale():
@@ -121,7 +127,7 @@ def test_zmax_floats():
 def test_zmin_zmax_range_color():
     img = img_gray / 100.0
     fig = px.imshow(img)
-    #assert not (fig.layout.coloraxis.cmin or fig.layout.coloraxis.cmax)
+    # assert not (fig.layout.coloraxis.cmin or fig.layout.coloraxis.cmax)
     fig1 = px.imshow(img, zmin=0.2, zmax=0.8)
     fig2 = px.imshow(img, range_color=[0.2, 0.8])
     assert fig1 == fig2
@@ -197,4 +203,41 @@ def test_imshow_source():
     assert np.all(decoded_img == img_rgb)
 
 
-# def test_imshow_source_dtype_zmax():
+@pytest.mark.parametrize(
+    "dtype",
+    [
+        np.uint8,
+        np.uint16,
+        np.int8,
+        np.int16,
+        np.int32,
+        np.int64,
+        np.float32,
+        np.float64,
+    ],
+)
+@pytest.mark.parametrize("contrast_rescaling", ["image", "dtype"])
+def test_imshow_source_dtype_zmax(dtype, contrast_rescaling):
+    img = np.arange(100, dtype=dtype).reshape((10, 10))
+    fig = px.imshow(img, use_binary_string=True, contrast_rescaling=contrast_rescaling)
+    if contrast_rescaling == "image":
+        assert (
+            np.max(
+                np.abs(
+                    rescale_intensity(img, in_range="image", out_range=np.uint8)
+                    - decode_image_string(fig.data[0].source)
+                )
+            )
+            < 1
+        )
+    else:
+        if dtype in [np.uint8, np.float32, np.float64]:
+            assert np.all(img == decode_image_string(fig.data[0].source))
+        else:
+            assert (
+                np.abs(
+                    np.max(decode_image_string(fig.data[0].source))
+                    - 255 * img.max() / np.iinfo(dtype).max
+                )
+                < 1
+            )