change(plot): add set_data support

src/aliby/utils/plot.py

@@ -13,13 +13,20 @@ from matplotlib import pyplot as plt
 def plot_overlay(
     bg: np.ndarray, fg: np.ndarray, alpha: float = 0.5, ax=plt
 ) -> None:
+    """
+    Plot two images, one on top of the other.
+    """
 
-    ax.imshow(bg, cmap=plt.cm.gray, interpolation="none")
-    ax.imshow(fg, alpha=alpha, interpolation="none")
-    ax.axis("off")
+    ax1 = ax.imshow(bg, cmap=plt.cm.gray, interpolation="none")
+    ax2 = ax.imshow(fg, alpha=alpha, interpolation="none")
+    plt.axis("off")
+    return ax1, ax2
 
 
 def plot_overlay_in_square(data: t.Tuple[np.ndarray, np.ndarray]):
+    """
+    Plot images in an automatically-arranged grid.
+    """
     specs = strategies.SquareStrategy("center").get_grid(len(data))
     for i, (gs, (tile, mask)) in enumerate(zip(specs, data)):
         ax = plt.subplot(gs)
@@ -27,6 +34,9 @@ def plot_overlay_in_square(data: t.Tuple[np.ndarray, np.ndarray]):
 
 
 def plot_in_square(data: t.Iterable):
+    """
+    Plot images in an automatically-arranged grid. Only takes one mask
+    """
     specs = strategies.SquareStrategy("center").get_grid(len(data))
     for i, (gs, datum) in enumerate(zip(specs, data)):
         ax = plt.subplot(gs)