Made fb.plot.images more customizable.

bethgelab · Feb 11, 2021 · dd3b180 · dd3b180
1 parent 2343d58
commit dd3b180
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 2 deletions.
diff --git a/foolbox/plot.py b/foolbox/plot.py
@@ -13,8 +13,10 @@ def images(
     nrows: Optional[int] = None,
     figsize: Optional[Tuple[float, float]] = None,
     scale: float = 1,
+    labels: Any = None,
+    return_fig: bool = False,
     **kwargs: Any,
-) -> None:
+) -> Optional[Tuple[Any, Any]]:
     import matplotlib.pyplot as plt
 
     x: ep.Tensor = ep.astensor(images)
@@ -57,7 +59,7 @@ def images(
         nrows=nrows,
         figsize=figsize,
         squeeze=False,
-        constrained_layout=True,
+        constrained_layout=False,
         **kwargs,
     )
 
@@ -68,5 +70,11 @@ def images(
             ax.set_yticks([])
             ax.axis("off")
             i = row * ncols + col
+            if labels is not None:
+                ax.set_title(labels[i])
             if i < len(x):
                 ax.imshow(x[i])
+
+    if return_fig:
+        return fig, axes
+    return None
diff --git a/tests/test_plot.py b/tests/test_plot.py
@@ -13,6 +13,10 @@ 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)
+    fbn.plot.images(
+        images, nrows=2, ncols=4, labels=[str(i) for i in range(len(images))]
+    )
+    fbn.plot.images(images, nrows=2, ncols=4, return_fig=True)
     with pytest.raises(ValueError):
         images = ep.zeros(dummy, (10, 3, 3, 3))
         fbn.plot.images(images)