ENH: Add auxiliary plotting to point selection

src/darsia/assistants/point_selection_assistant.py

@@ -4,6 +4,8 @@
 from warnings import warn
 
 import matplotlib.pyplot as plt
+import numpy as np
+import skimage
 
 import darsia
 
@@ -30,6 +32,20 @@ def __init__(self, img: darsia.Image, **kwargs) -> None:
 
         # Output mode
         self.return_coordinates = kwargs.get("coordinates", False)
+        """Flag controlling whether to return coordinates or voxels."""
+
+        # Activate masking if labels are provided
+        self.labels: Optional[darsia.Image] = kwargs.get("labels", None)
+        """Labels for the image."""
+        if self.labels is not None:
+            # Expand image with a (active) mask
+            coarse_img = skimage.img_as_float(
+                darsia.resize(
+                    self.img, fx=0.1, fy=0.1, interpolation="inter_nearest"
+                ).img
+            )
+            mask = np.ones(coarse_img.shape[:2], dtype=np.float32)
+            self.coarse_masked_img = np.dstack((coarse_img, mask))
 
     def __call__(self) -> Union[darsia.CoordinateArray, darsia.VoxelArray]:
         """Call the assistant."""
@@ -120,6 +136,23 @@ def _on_mouse_click(self, event: Any) -> None:
                     "go",
                     markersize=10,
                 )
+
+                # Auxiliary plot: Mark labeled region
+                if self.labels is not None:
+                    label = self.labels.img[int(event.ydata), int(event.xdata)]
+                    mask = (
+                        darsia.resize(
+                            skimage.img_as_float(self.labels.img == label),
+                            shape=self.coarse_masked_img.shape[:2],
+                            interpolation="inter_nearest",
+                        )
+                        > 0.5
+                    )
+                    self.coarse_masked_img[mask, -1] = 0.3
+                    plt.figure("Auxiliary: Masked image")
+                    plt.imshow(self.coarse_masked_img)
+                    plt.show()
+
                 self.fig.canvas.draw()
                 self._print_current_selection()