Allow normalising RGB images

Zac-HD · Zac-HD · commit f0950acb5b95 · 2018-01-11T00:45:32.000+11:00
diff --git a/doc/users/credits.rst b/doc/users/credits.rst
@@ -386,6 +386,7 @@ Yu Feng,
 Yunfei Yang,
 Yuri D'Elia,
 Yuval Langer,
+Zac Hatfield-Dodds,
 Zach Pincus,
 Zair Mubashar,
 alex,
diff --git a/doc/users/next_whats_new/2018_01_11_imshow-rgb-scaling-clipping.rst b/doc/users/next_whats_new/2018_01_11_imshow-rgb-scaling-clipping.rst
@@ -0,0 +1,6 @@
+`Axes.imshow` can normalise RGB images
+--------------------------------------
+
+`Axes.imshow` can now apply the ``norm`` or ``vmin`` and ``vmax`` arguments
+to RGB images, allowing easy display of high-bit-depth data, or adjustments
+of brightness and contrast.
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -5222,18 +5222,22 @@ def imshow(self, X, cmap=None, norm=None, aspect=None,
         Parameters
         ----------
         X : array_like, shape (n, m) or (n, m, 3) or (n, m, 4)
-            Display the image in `X` to current axes.  `X` may be an
-            array or a PIL image. If `X` is an array, it
-            can have the following shapes and types:
+            Display the image in `X` to current axes.  `X` may be an array
+            or a PIL image. If `X` is an array, it may have the following
+            shapes, and either floating or integer dtypes:
 
-            - MxN -- values to be mapped (float or int)
-            - MxNx3 -- RGB (float or uint8)
-            - MxNx4 -- RGBA (float or uint8)
+            - MxN -- values to be mapped
+            - MxNx3 -- RGB
+            - MxNx4 -- RGBA
 
-            The value for each component of MxNx3 and MxNx4 float arrays
-            should be in the range 0.0 to 1.0. MxN arrays are mapped
-            to colors based on the `norm` (mapping scalar to scalar)
-            and the `cmap` (mapping the normed scalar to a color).
+            MxN arrays are mapped to colors based on the `norm` (mapping
+            scalar to scalar) and the `cmap` (mapping the normed scalar to
+            a color).
+
+            Elements of RGB and RGBA arrays represent pixels of an MxN image.
+            All values - unless normalised with `norm` or `vmin` and `vmax` -
+            should be in the range [0 .. 1] for floats or [0 .. 255] for
+            integers.
 
         cmap : `~matplotlib.colors.Colormap`, optional, default: None
             If None, default to rc `image.cmap` value. `cmap` is ignored
@@ -5271,11 +5275,13 @@ def imshow(self, X, cmap=None, norm=None, aspect=None,
 
         vmin, vmax : scalar, optional, default: None
             `vmin` and `vmax` are used in conjunction with norm to normalize
-            luminance data.  Note if you pass a `norm` instance, your
-            settings for `vmin` and `vmax` will be ignored.
+            luminance, RGB, or RGBA data.  Note if you pass a `norm` instance,
+            your settings for `vmin` and `vmax` will be ignored.
 
         alpha : scalar, optional, default: None
-            The alpha blending value, between 0 (transparent) and 1 (opaque)
+            The alpha blending value, between 0 (transparent) and 1 (opaque).
+            For RGBA images, this is applied in addition to the data's alpha
+            channel.
 
         origin : ['upper' | 'lower'], optional, default: None
             Place the [0,0] index of the array in the upper left or lower left
@@ -5341,6 +5347,21 @@ def imshow(self, X, cmap=None, norm=None, aspect=None,
         if aspect is None:
             aspect = rcParams['image.aspect']
         self.set_aspect(aspect)
+
+        if X.ndim == 3:
+            # Much of the standard machinery doesn't work implicitly for RGB
+            # and RGBA images, so we replicate the effects here.
+
+            # RGB and RGBA arrays can't be implicitly normalised, but explicit
+            # normalisation works just fine!  However, we only normalise if
+            # an argument is given - otherwise you can't display a dark image.
+            if norm is None and (vmin is not None or vmax is not None):
+                norm = mcolors.Normalize(vmin, vmax, clip=True)
+            if norm is not None:
+                X = norm(X)
+                # Disable later styling based on these values
+                norm, vmin, vmax = None, None, None
+
         im = mimage.AxesImage(self, cmap, norm, interpolation, origin, extent,
                               filternorm=filternorm, filterrad=filterrad,
                               resample=resample, **kwargs)
