tweaks so that zmin/zmax are set less often

Author: emmanuelle

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

@@ -57,7 +57,9 @@ def _array_to_b64str(img, backend="pil", compression=4):
         sh = img.shape
         if ndim == 3:
             img = img.reshape((sh[0], sh[1] * sh[2]))
-        w = png.Writer(sh[1], sh[0], greyscale=(ndim == 2), alpha=alpha, compression=compression)
+        w = png.Writer(
+            sh[1], sh[0], greyscale=(ndim == 2), alpha=alpha, compression=compression
+        )
         img_png = png.from_array(img, mode=mode)
         prefix = "data:image/png;base64,"
         with BytesIO() as stream:
@@ -81,11 +83,11 @@ def _vectorize_zvalue(z):
     if z is None:
         return z
     elif np.isscalar(z):
-        return [z] * 3 + [1]
+        return [z] * 3 + [255]
     elif len(z) == 1:
-        return list(z) * 3 + [1]
+        return list(z) * 3 + [255]
     elif len(z) == 3:
-        return list(z) + [1]
+        return list(z) + [255]
     elif len(z) == 4:
         return z
     else:
@@ -311,11 +313,7 @@ def imshow(
             raise ValueError("Binary strings cannot be used with pandas arrays")
         has_nans = True
     else:
-        has_nans = np.any(np.isnan(img))
-        if has_nans and binary_string:
-            raise ValueError(
-                "Binary strings cannot be used with arrays containing NaNs"
-            )
+        has_nans = False
 
     # --------------- Starting from here img is always a numpy array --------
     img = np.asanyarray(img)
@@ -340,9 +338,9 @@ def imshow(
         if (zmax is not None or binary_string) and zmin is None:
             zmin = img.min()
     else:
-        if zmax is None and (img.dtype is not np.uint8 or img.ndim == 2):
+        if zmax is None and (img.dtype != np.uint8 or img.ndim == 2):
             zmax = _infer_zmax_from_type(img)
-        if zmin is None:
+        if zmin is None and zmax is not None:
             zmin = 0
 
         # For 2d data, use Heatmap trace, unless binary_string is True
@@ -377,9 +375,12 @@ def imshow(
 
     # For 2D+RGB data, use Image trace
     elif img.ndim == 3 and img.shape[-1] in [3, 4] or (img.ndim == 2 and binary_string):
-        zmin, zmax = _vectorize_zvalue(zmin), _vectorize_zvalue(zmax)
+        if zmin is not None and zmax is not None:
+            zmin, zmax = _vectorize_zvalue(zmin), _vectorize_zvalue(zmax)
         if binary_string:
-            if img.ndim == 2:
+            if zmin is None and zmax is None:
+                img_rescaled = img
+            elif img.ndim == 2:
                 img_rescaled = rescale_intensity(
                     img, in_range=(zmin[0], zmax[0]), out_range=np.uint8
                 )
@@ -403,7 +404,8 @@ def imshow(
             )
             trace = go.Image(source=img_str)
         else:
-            trace = go.Image(z=img, zmin=zmin, zmax=zmax)
+            colormodel = "rgb" if img.shape[-1] == 3 else "rgba"
+            trace = go.Image(z=img, zmin=zmin, zmax=zmax, colormodel=colormodel)
         layout = {}
         if origin == "lower":
             layout["yaxis"] = dict(autorange=True)

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

@@ -26,10 +26,7 @@ def decode_image_string(image_string):
 @pytest.mark.parametrize("binary_string", [False, True])
 def test_rgb_uint8(binary_string):
     fig = px.imshow(img_rgb, binary_string=binary_string)
-    if not binary_string:
-        assert fig.data[0]["zmax"] == (255, 255, 255, 1)
-    else:
-        assert fig.data[0]["zmax"] is None
+    assert fig.data[0]["zmax"] is None
 
 
 def test_zmax():
@@ -39,10 +36,10 @@ def test_zmax():
         (100,),
         [100, 100, 100],
         (100, 100, 100),
-        (100, 100, 100, 1),
+        (100, 100, 100, 255),
     ]:
         fig = px.imshow(img_rgb, zmax=zmax, binary_string=False)
-        assert fig.data[0]["zmax"] == (100, 100, 100, 1)
+        assert fig.data[0]["zmax"] == (100, 100, 100, 255)
 
 
 def test_automatic_zmax_from_dtype():
@@ -56,7 +53,11 @@ def test_automatic_zmax_from_dtype():
         img = np.array([0, 1], dtype=key)
         img = np.dstack((img,) * 3)
         fig = px.imshow(img, binary_string=False)
-        assert fig.data[0]["zmax"] == (val, val, val, 1)
+        # For uint8 in "infer" mode we don't pass zmin/zmax unless specified
+        if key in [np.uint8, np.bool]:
+            assert fig.data[0]["zmax"] is None
+        else:
+            assert fig.data[0]["zmax"] == (val, val, val, 255)
 
 
 @pytest.mark.parametrize("binary_string", [False, True])
@@ -91,15 +92,23 @@ def test_wrong_dimensions():
             _ = px.imshow(img)
 
 
-def test_nan_inf_data():
+@pytest.mark.parametrize("binary_string", [False, True])
+def test_nan_inf_data(binary_string):
     imgs = [np.ones((20, 20)), 255 * np.ones((20, 20))]
     zmaxs = [1, 255]
     for zmax, img in zip(zmaxs, imgs):
         img[0] = 0
         img[10:12] = np.nan
         # the case of 2d/heatmap is handled gracefully by the JS trace but I don't know how to check it
-        fig = px.imshow(np.dstack((img,) * 3))
-        assert fig.data[0]["zmax"] == (zmax, zmax, zmax, 1)
+        fig = px.imshow(
+            np.dstack((img,) * 3),
+            binary_string=binary_string,
+            contrast_rescaling="minxmax",
+        )
+        if not binary_string:
+            assert fig.data[0]["zmax"] == (zmax, zmax, zmax, 255)
+        else:
+            assert fig.data[0]["zmax"] is None
 
 
 def test_zmax_floats():
@@ -113,7 +122,7 @@ def test_zmax_floats():
     zmaxs = [1, 1, 255, 65535]
     for zmax, img in zip(zmaxs, imgs):
         fig = px.imshow(img, binary_string=False)
-        assert fig.data[0]["zmax"] == (zmax, zmax, zmax, 1)
+        assert fig.data[0]["zmax"] == (zmax, zmax, zmax, 255)
     # single-channel
     imgs = [
         np.ones((5, 5)),