Support 2D and RGB images for histograms

Author: dstansby

CHANGELOG.rst

@@ -0,0 +1,12 @@
+0.0.2
+=====
+
+New features
+------------
+- `HistogramWidget` now shows individual histograms for RGB channels when
+  present.
+
+
+Bug fixes
+---------
+- `HistogramWidget` now works properly with 2D images.

src/napari_matplotlib/histogram.py

@@ -6,6 +6,9 @@
 __all__ = ["HistogramWidget"]
 
 
+_COLORS = {"r": "tab:red", "g": "tab:green", "b": "tab:blue"}
+
+
 class HistogramWidget(NapariMPLWidget):
     """
     Widget to display a histogram of the currently selected layer.
@@ -40,9 +43,27 @@ def hist_current_layer(self) -> None:
         """
         self.axes.clear()
         layer = self.layer
-        z = self.viewer.dims.current_step[0]
         bins = np.linspace(np.min(layer.data), np.max(layer.data), 100)
-        data = layer.data[z]
-        self.axes.hist(data.ravel(), bins=bins)
-        self.axes.set_title(f"{layer.name}, z={z}")
+
+        if layer.data.ndim - layer.rgb == 3:
+            # 3D data, can be single channel or RGB
+            data = layer.data[self.current_z]
+            self.axes.set_title(f"z={self.current_z}")
+        else:
+            data = layer.data
+
+        if layer.rgb:
+            # Histogram RGB channels independently
+            for i, c in enumerate("rgb"):
+                self.axes.hist(
+                    data[..., i].ravel(),
+                    bins=bins,
+                    label=c,
+                    histtype="step",
+                    color=_COLORS[c],
+                )
+        else:
+            self.axes.hist(data.ravel(), bins=bins, label=layer.name)
+
+        self.axes.legend()
         self.canvas.draw()

src/napari_matplotlib/tests/test_histogram.py

@@ -1,10 +1,15 @@
 import numpy as np
+import pytest
 
 from napari_matplotlib import HistogramWidget
 
 
-def test_example_q_widget(make_napari_viewer):
+# Test with 2D and 3D data
+@pytest.mark.parametrize(
+    "data", [np.random.random((100, 100)), np.random.random((100, 100, 100))]
+)
+def test_example_q_widget(make_napari_viewer, data):
     # Smoke test adding a histogram widget
     viewer = make_napari_viewer()
-    viewer.add_image(np.random.random((100, 100)))
+    viewer.add_image(data)
     HistogramWidget(viewer)