Fix single channel plotting (#631)

* Fix single channel plotting * Improve test_plot covering 1 and 3 channel images
bethgelab · Jun 4, 2021 · cc7499c · cc7499c
1 parent cbad38f
commit cc7499c
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/foolbox/plot.py b/foolbox/plot.py
@@ -69,4 +69,7 @@ def images(
             ax.axis("off")
             i = row * ncols + col
             if i < len(x):
-                ax.imshow(x[i])
+                if x.shape[-1] == 1:
+                    ax.imshow(x[i][:, :, 0])
+                else:
+                    ax.imshow(x[i])
diff --git a/tests/test_plot.py b/tests/test_plot.py
@@ -13,6 +13,9 @@ def test_plot(dummy: ep.Tensor) -> None:
     fbn.plot.images(images, ncols=3)
     fbn.plot.images(images, nrows=2, ncols=6)
     fbn.plot.images(images, nrows=2, ncols=4)
+    # test for single channel images
+    images = ep.zeros(dummy, (10, 32, 32, 1))
+    fbn.plot.images(images)
     with pytest.raises(ValueError):
         images = ep.zeros(dummy, (10, 3, 3, 3))
         fbn.plot.images(images)