Fix matplotlib.colormaps

google · Jan 4, 2023 · a0f0583 · a0f0583
1 parent 01fdedc
commit a0f0583
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 4 deletions.
diff --git a/mediapy/__init__.py b/mediapy/__init__.py
@@ -97,7 +97,7 @@
 """
 
 __docformat__ = 'google'
-__version__ = '1.1.3'
+__version__ = '1.1.4'
 __version_info__ = tuple(int(num) for num in __version__.split('.'))
 
 import base64
@@ -124,7 +124,6 @@
 
 import IPython.display
 import matplotlib
-import matplotlib.pyplot as plt
 import numpy as np
 import numpy.typing as npt
 import PIL.Image
@@ -756,7 +755,10 @@ def to_rgb(
   vmax = np.amax(np.where(np.isfinite(a), a, -np.inf)) if vmax is None else vmax
   a = (a.astype('float') - vmin) / (vmax - vmin + np.finfo(float).eps)
   if isinstance(cmap, str):
-    rgb_from_scalar = matplotlib.colormaps[cmap]
+    if hasattr(matplotlib, 'colormaps'):
+      rgb_from_scalar = matplotlib.colormaps[cmap]  # Newer version.
+    else:
+      rgb_from_scalar = matplotlib.pyplot.cm.get_cmap(cmap)
   else:
     rgb_from_scalar = cmap
   a = rgb_from_scalar(a)

diff --git a/mediapy_test.py b/mediapy_test.py
@@ -394,7 +394,11 @@ def test_uint8_to_rgb(self):
     a = np.array([100, 120, 140], dtype=np.uint8)
 
     def gray(x):
-      return matplotlib.colormaps['gray'](x)[..., :3]
+      if hasattr(matplotlib, 'colormaps'):
+        cmap = matplotlib.colormaps['gray']  # Newer version.
+      else:
+        cmap = matplotlib.pyplot.cm.get_cmap('gray')
+      return cmap(x)[..., :3]
 
     self.assert_all_close(media.to_rgb(a), gray([0.0, 0.5, 1.0]))
     self.assert_all_close(