Make matplotlib simulation and daltonization work again

.gitignore

@@ -7,6 +7,7 @@ dist
 .ipynb_checkpoints/
 doc/colourmaps.pdf
 .vscode
-.coverage
+.coverage*
 .tox
 coverage.json
+result-images/

README.md

@@ -2,15 +2,15 @@
 
 ![https://github.com/joergdietrich/daltonize/actions](https://img.shields.io/github/actions/workflow/status/joergdietrich/daltonize/main.yml)  ![](https://img.shields.io/endpoint?url=https://gist.githubusercontent.com/joergdietrich/9deb619232c8098b5e15d259ef5ed534/raw/covbadge.json)
 
-Daltonize simulates the three types of dichromatic color blindness and
+Daltonize simulates the three types of dichromatic color blindness for
 images and matplotlib figures. Generalizing and omitting a lot of
 details these types are:
 
 * Deuteranopia: green weakness
 * Protanopia: red weakness
 * Tritanopia: blue weakness (extremely rare)
 
-Daltonize can also adjust the color palette of an input image such
+Daltonize can also adjust the color palette of an input image or matplotlib figure such
 that a color blind person can perceive the full information
 content. It can be used as a command line tool to convert pixel images
 but also as a Python module. If used as the latter, it provides an API
@@ -19,6 +19,10 @@ to simulate and correct for color blindness in matplotlib figures.
 This allows to create color blind friendly vector graphics suitable
 for publication.
 
+Color vision deficiencies are in fact very complex and can differ in intensity from person to person.
+The algorithms used in here and in many other comparable software packages are based on often simplifying assumptions. [Nicolas Burrus](http://nicolas.burrus.name/) discusses these simplification and reviews daltonize and other software packages in this [blog post](https://daltonlens.org/opensource-cvd-simulation/).
+
+
 ## Installation
 
 ```
@@ -97,15 +101,15 @@ the `-s/--simulate` option and `-t/--type d` or `p` on these images.
 ### Deuteranopia
 
 ```
-python daltonize.py -s -t=d example_images/Ishihara_Plate_3.jpg example_images/Ishihara_Plate_3-Deuteranopia.jpg
+daltonize -s -t=d example_images/Ishihara_Plate_3.jpg example_images/Ishihara_Plate_3-Deuteranopia.jpg
 ```
 
 ![IshiharaPlate3](example_images/Ishihara_Plate_3-Deuteranopia.jpg)
 
 ### Protanopia
 
 ```
-python daltonize.py -s -t=p example_images/Ishihara_Plate_3.jpg example_images/Ishihara_Plate_3-Protanopia.jpg
+daltonize -s -t=p example_images/Ishihara_Plate_3.jpg example_images/Ishihara_Plate_3-Protanopia.jpg
 ```
 
 ![IshiharaPlate3](example_images/Ishihara_Plate_3-Protanopia.jpg)

daltonize/daltonize.py

@@ -21,6 +21,7 @@
 except ImportError:
     _NO_MPL = True
 
+COLOR_KEYS = ("color", "fc", "ec", "mec", "mfc", "mfcalt", "cmap", "array")
 
 def transform_colorspace(img, mat):
     """Transform image to a different color space.
@@ -44,7 +45,7 @@ def simulate(rgb, color_deficit="d"):
     Arguments:
     ----------
     rgb : array of shape (M, N, 3)
-        original image in RGB format
+        original image in RGB format, values must be [0, 1] bounded
     color_deficit : {"d", "p", "t"}, optional
         type of colorblindness, d for deuteronopia (default),
         p for protonapia,
@@ -74,7 +75,7 @@ def simulate(rgb, color_deficit="d"):
     sim_lms = transform_colorspace(lms, cb_matrices[color_deficit])
     # Transform back to RBG
     sim_rgb = transform_colorspace(sim_lms, lms2rgb)
-    return sim_rgb
+    return np.clip(sim_rgb, 0, 1)
 
 
 def daltonize(rgb, color_deficit='d'):
@@ -84,7 +85,7 @@ def daltonize(rgb, color_deficit='d'):
     Arguments:
     ----------
     rgb : array of shape (M, N, 3)
-        original image in RGB format
+        original image in RGB format, values must be [0, 1] bounded
     color_deficit : {"d", "p", "t"}, optional
         type of colorblindness, d for deuteronopia (default),
         p for protonapia,
@@ -102,7 +103,7 @@ def daltonize(rgb, color_deficit='d'):
     # they can see.
     err = transform_colorspace(rgb - sim_rgb, err2mod)
     dtpn = err + rgb
-    return dtpn
+    return np.clip(dtpn, 0, 1)
 
 
 def array_to_img(arr, gamma=2.4):
@@ -120,34 +121,12 @@ def array_to_img(arr, gamma=2.4):
     """
     # clip values to lie in the range [0, 255]
     arr = inverse_gamma_correction(arr, gamma=gamma)
-    arr = clip_array(arr)
+    arr = np.clip(arr, 0, 255)
     arr = arr.astype('uint8')
     img = Image.fromarray(arr, mode='RGB')
     return img
 
 
-def clip_array(arr, min_value=0, max_value=255):
-    """Ensure that all values in an array are between min and max values.
-
-    Arguments:
-    ----------
-    arr : array_like
-    min_value : float, optional
-        default 0
-    max_value : float, optional
-        default 255
-
-    Returns:
-    --------
-    arr : array_like
-        clipped such that all values are min_value <= arr <= max_value
-    """
-    comp_arr = np.ones_like(arr)
-    arr = np.maximum(comp_arr * min_value, arr)
-    arr = np.minimum(comp_arr * max_value, arr)
-    return arr
-
-
 def get_child_colors(child, mpl_colors):
     """
     Recursively enter all colors of a matplotlib objects and its
@@ -232,30 +211,28 @@ def get_key_colors(mpl_colors, rgb, alpha):
     if _NO_MPL is True:
         raise ImportError("matplotlib not found, "
                           "can only deal with pixel images")
-    cc = mpl.colors.ColorConverter()  # pylint: disable=invalid-name
-    # Note that the order must match the insertion order in
-    # get_child_colors()
-    color_keys = ("color", "fc", "ec", "mec", "mfc", "mfcalt", "cmap", "array")
-    for color_key in color_keys:
+    for color_key in COLOR_KEYS:
         try:
             color = mpl_colors[color_key]
             # skip unset colors, otherwise they are turned into black.
             if isinstance(color, str) and color == 'none':
                 continue
+            if isinstance(color, mpl.colors.ListedColormap):
+                continue
             if isinstance(color, mpl.colors.LinearSegmentedColormap):
                 rgba = color(np.arange(color.N))
             elif isinstance(color, np.ndarray) and color_key == "array":
                 color = color.reshape(-1, 3) / 255
                 a = np.zeros((color.shape[0], 1))  # pylint: disable=invalid-name
                 rgba = np.hstack((color, a))
             else:
-                rgba = cc.to_rgba_array(color)
+                rgba = mpl.colors.to_rgba_array(color)
             rgb = np.append(rgb, rgba[:, :3])
             alpha = np.append(alpha, rgba[:, 3])
         except KeyError:
             pass
         for key in mpl_colors.keys():
-            if key in color_keys:
+            if key in COLOR_KEYS:
                 continue
             rgb, alpha = get_key_colors(mpl_colors[key], rgb, alpha)
     return rgb, alpha
@@ -293,10 +270,8 @@ def _set_colors_from_array(instance, mpl_colors, rgba, i=0):
     Set object instance colors to the modified ones in rgba.
     """
     cc = mpl.colors.ColorConverter()  # pylint: disable=invalid-name
-    # Note that the order must match the insertion order in
-    # get_child_colors()
-    color_keys = ("color", "fc", "ec", "mec", "mfc", "mfcalt", "cmap", "array")
-    for color_key in color_keys:
+
+    for color_key in COLOR_KEYS:
         try:
             color = mpl_colors[color_key]
             if isinstance(color, mpl.colors.LinearSegmentedColormap):
@@ -307,6 +282,8 @@ def _set_colors_from_array(instance, mpl_colors, rgba, i=0):
                 # skip unset colors, otherwise they are turned into black.
                 if isinstance(color, str) and color == 'none':
                     continue
+                if isinstance(color, mpl.colors.ListedColormap):
+                    continue
                 color_shape = cc.to_rgba_array(color).shape
                 j = color_shape[0]
             target_color = rgba[i: i + j, :]
@@ -316,6 +293,10 @@ def _set_colors_from_array(instance, mpl_colors, rgba, i=0):
             if color_key == "color":
                 instance.set_color(target_color)
             elif color_key == "fc":
+                # Circumvent https://github.com/matplotlib/matplotlib/issues/13131 by setting
+                # the cmap to None, which would otherwise take precedence over explicitly set colors.
+                if hasattr(instance, "set_array"):
+                    instance.set_array(None)
                 instance.set_facecolor(target_color)
             elif color_key == "ec":
                 instance.set_edgecolor(target_color)
@@ -371,7 +352,6 @@ def _join_rgb_alpha(rgb, alpha):
     """
     Combine (m, n, 3) rgb and (m, n) alpha array into (m, n, 4) rgba.
     """
-    rgb = clip_array(rgb, 0, 1)
     r, g, b = np.split(rgb, 3, 2)  # pylint: disable=invalid-name, unbalanced-tuple-unpacking
     rgba = np.concatenate((r, g, b, alpha.reshape(alpha.size, 1, 1)),
                           axis=2).reshape(-1, 4)
@@ -397,13 +377,17 @@ def simulate_mpl(fig, color_deficit='d', copy=False):
     --------
     fig : matplotlib.figure.Figure
     """
+    # Ensure that everything that was specified is actually drawn
+    fig.canvas.draw()
     if copy:
         # mpl.transforms cannot be copy.deepcopy()ed. Thus we resort
         # to pickling.
         pfig = pickle.dumps(fig)
         fig = pickle.loads(pfig)
     rgb, alpha, mpl_colors = _prepare_for_transform(fig)
-    sim_rgb = simulate(array_to_img(rgb * 255), color_deficit) / 255
+    rgb = inverse_gamma_correction(rgb) / 255
+    sim_rgb = simulate(rgb, color_deficit)
+    sim_rgb = gamma_correction(sim_rgb * 255)
     rgba = _join_rgb_alpha(sim_rgb, alpha)
     set_mpl_colors(mpl_colors, rgba)
     fig.canvas.draw()
@@ -429,13 +413,17 @@ def daltonize_mpl(fig, color_deficit='d', copy=False):
     --------
     fig : matplotlib.figure.Figure
     """
+    # Ensure that everything that was specified is actually drawn
+    fig.canvas.draw()
     if copy:
         # mpl.transforms cannot be copy.deepcopy()ed. Thus we resort
         # to pickling.
         pfig = pickle.dumps(fig)
         fig = pickle.loads(pfig)
     rgb, alpha, mpl_colors = _prepare_for_transform(fig)
-    dtpn = daltonize(array_to_img(rgb * 255), color_deficit) / 255
+    rgb = inverse_gamma_correction(rgb) / 255
+    dtpn = daltonize(rgb, color_deficit)
+    sim_rgb = gamma_correction(dtpn * 255)
     rgba = _join_rgb_alpha(dtpn, alpha)
     set_mpl_colors(mpl_colors, rgba)
     fig.canvas.draw()
@@ -444,9 +432,9 @@ def daltonize_mpl(fig, color_deficit='d', copy=False):
 def gamma_correction(rgb, gamma=2.4):
     """
     Apply sRGB gamma correction
-    :param rgb:
+    :param rgb: array of shape (M, N, 3) with sRGB values in the range [0, 255]
     :param gamma:
-    :return: linear_rgb
+    :return: linear_rgb, array of shape (M, N, 3) with linear sRGB values between in the range [0, 1]
     """
     linear_rgb = np.zeros_like(rgb, dtype=np.float16)
     for i in range(3):
@@ -461,7 +449,7 @@ def inverse_gamma_correction(linear_rgb, gamma=2.4):
 
     :param linear_rgb: array of shape (M, N, 3) with linear sRGB values between in the range [0, 1]
     :param gamma: float
-    :return: array of shape (M, N, 3) with inverse gamma correction applied
+    :return: array of shape (M, N, 3) with inverse gamma correction applied, values in the range [0, 255]
     """
     rgb = np.zeros_like(linear_rgb, dtype=np.float16)
     for i in range(3):
@@ -490,7 +478,7 @@ def main():
                              "(most common)")
     parser.add_argument("-g", "--gamma", type=float, default=2.4,
                         help="value of gamma correction to be applied before transformation. The default "
-                             "applies the standard sRGB correction with an exponent of 2.4"
+                             "applies the standard sRGB correction with an exponent of 2.4 "
                              "Use 1 for no gamma correction.")
     args = parser.parse_args()
 

daltonize/tests/test_daltonize.py

@@ -5,34 +5,79 @@
 from PIL import Image
 
 import numpy as np
+from matplotlib import pyplot as plt
+from matplotlib.testing.decorators import image_comparison
+from matplotlib.testing.conftest import mpl_test_settings
 from numpy.testing import (assert_equal, assert_array_almost_equal)
 
-from daltonize.daltonize import gamma_correction, inverse_gamma_correction, simulate, array_to_img
-
-def test_gamma_correction():
-    rgb = np.array([[[0, 10, 11, 25, 128, 255]]]).reshape((-1, 1, 3))
-    expected = np.array([[[0.      , 0.003035, 0.003347]],
-                         [[0.00972 , 0.2158  , 1.      ]]], dtype=np.float16)
-    linear_rgb = gamma_correction(rgb)
-    assert_array_almost_equal(linear_rgb, expected)
-    expected = np.array([[[0.      , 0.003035, 0.093   ]],
-                         [[0.145   , 0.528   , 1.      ]]], dtype=np.float16)
-    linear_rgb = gamma_correction(rgb, gamma=1)
-    assert_array_almost_equal(linear_rgb, expected)
-
-def test_inverse_gamma_correction():
-    rgb = np.array([[[0, 10, 11, 25, 128, 255]]]).reshape((-1, 1, 3))
-    assert_array_almost_equal(rgb, inverse_gamma_correction(gamma_correction(rgb)))
-
-@pytest.mark.parametrize("type, ref_img_path", [("d", Path("daltonize/tests/data/colored_crayons_d.png")),
-                                                ("p", Path("daltonize/tests/data/colored_crayons_p.png")),
-                                                ("t", Path("daltonize/tests/data/colored_crayons_t.png"))])
-def test_simulation(type, ref_img_path):
-    gamma = 2.4
-    orig_img_path = Path("daltonize/tests/data/colored_crayons.png")
-    orig_img = np.asarray(Image.open(orig_img_path).convert("RGB"), dtype=np.float16)
-    orig_img = gamma_correction(orig_img, gamma)
-    simul_rgb = simulate(orig_img, type)
-    simul_img = np.asarray(array_to_img(simul_rgb, gamma=gamma))
-    ref_img = np.asarray(Image.open(ref_img_path).convert("RGB"))
-    assert_array_almost_equal(simul_img, ref_img)
+from daltonize.daltonize import *
+
+class TestDaltonize():
+    def test_gamma_correction(self):
+        rgb = np.array([[[0, 10, 11, 25, 128, 255]]]).reshape((-1, 1, 3))
+        expected = np.array([[[0.      , 0.003035, 0.003347]],
+                            [[0.00972 , 0.2158  , 1.      ]]], dtype=np.float16)
+        linear_rgb = gamma_correction(rgb)
+        assert_array_almost_equal(linear_rgb, expected)
+        expected = np.array([[[0.      , 0.003035, 0.093   ]],
+                            [[0.145   , 0.528   , 1.      ]]], dtype=np.float16)
+        linear_rgb = gamma_correction(rgb, gamma=1)
+        assert_array_almost_equal(linear_rgb, expected)
+
+    def test_inverse_gamma_correction(self):
+        rgb = np.array([[[0, 10, 11, 25, 128, 255]]]).reshape((-1, 1, 3))
+        assert_array_almost_equal(rgb, inverse_gamma_correction(gamma_correction(rgb)))
+
+    @pytest.mark.parametrize("type, ref_img_path", [("d", Path("daltonize/tests/data/colored_crayons_d.png")),
+                                                    ("p", Path("daltonize/tests/data/colored_crayons_p.png")),
+                                                    ("t", Path("daltonize/tests/data/colored_crayons_t.png"))])
+    def test_simulate(self, type, ref_img_path):
+        gamma = 2.4
+        orig_img_path = Path("daltonize/tests/data/colored_crayons.png")
+        orig_img = np.asarray(Image.open(orig_img_path).convert("RGB"), dtype=np.float16)
+        orig_img = gamma_correction(orig_img, gamma)
+        simul_rgb = simulate(orig_img, type)
+        simul_img = np.asarray(array_to_img(simul_rgb, gamma=gamma))
+        ref_img = np.asarray(Image.open(ref_img_path).convert("RGB"))
+        assert_array_almost_equal(simul_img, ref_img)
+
+    @pytest.mark.parametrize("type, ref_img_path", [("d", Path("daltonize/tests/data/colored_crayons_daltonized_d.png")),
+                                                    ("p", Path("daltonize/tests/data/colored_crayons_daltonized_p.png")),
+                                                    ("t", Path("daltonize/tests/data/colored_crayons_daltonized_t.png"))])
+    def test_daltonize(self, type, ref_img_path):
+        gamma = 2.4
+        orig_img_path = Path("daltonize/tests/data/colored_crayons.png")
+        orig_img = np.asarray(Image.open(orig_img_path).convert("RGB"), dtype=np.float16)
+        orig_img = gamma_correction(orig_img, gamma)
+        simul_rgb = daltonize(orig_img, type)
+        simul_img = np.asarray(array_to_img(simul_rgb, gamma=gamma))
+        ref_img = np.asarray(Image.open(ref_img_path).convert("RGB"))
+        assert_array_almost_equal(simul_img, ref_img)
+
+    @image_comparison([
+        "imshow_crayons_d_daltonized.png",
+        "imshow_crayons_d.png", 
+        ], remove_text=True,
+                       savefig_kwarg={"dpi": 40})
+    def test_mpl_imshow(self):
+        fig = plt.figure()
+        img = Image.open(Path("daltonize/tests/data/colored_crayons.png"))
+        plt.imshow(img)
+        sim_fig = simulate_mpl(fig, copy=True)
+        fig = daltonize_mpl(fig)
+
+    @image_comparison([
+        "imshow_scatter_with_colorbar_p_daltonized.png",
+        "imshow_scatter_with_colorbar_p.png", 
+        ], remove_text=True,
+                       savefig_kwarg={"dpi": 40})
+    def test_mpl_scatter_with_colorbar(self):
+        fig = plt.figure()
+        np.random.seed(1)
+        x = np.random.random(size=50)
+        y = np.random.random(size=50)
+        plt.scatter(x, y, c=y, cmap="rainbow")
+        plt.colorbar()
+        sim_fig = simulate_mpl(fig, color_deficit="p", copy=True)
+        fig = daltonize_mpl(fig, color_deficit="p")
+

requirements.txt

@@ -1,2 +1,3 @@
+matplotlib>=3.0.0
 numpy>=1.9.0
 Pillow