new mask drawing function using matplotlib

pyproject.toml

@@ -55,7 +55,9 @@ ui = [
     "holoviews",
     "hvplot",
     "ipyaggrid",
-    "matplotlib"
+    "matplotlib",
+    "ipympl",
+    "ipywidgets"
 ]
 doc = [
         "sphinx",

src/PyHyperScattering/IntegrationUtils.py

@@ -7,10 +7,13 @@
 try:
     import matplotlib.pyplot as plt
     from matplotlib.colors import LogNorm,Normalize
+    from matplotlib.path import Path as MplPath
     import holoviews as hv
     import hvplot.xarray
     import skimage.draw
-
+    import ipywidgets as widgets
+    from IPython.display import display
+
 except (ModuleNotFoundError,ImportError):
     warnings.warn('Could not import a dependency for interactive integration utils.  Install pyhyperscattering[ui] or pyhyperscattering[all].',stacklevel=2)
 import pandas as pd
@@ -206,6 +209,179 @@ def mask(self):
         return mask
 
 
+class DrawMaskMatplotlib:
+    '''
+    Interactive polygon mask tool built on matplotlib/ipympl, as an alternative to DrawMask
+    for notebook environments where DrawMask's holoviews/bokeh-based .ui() does not render
+    as an interactive widget.
+
+    Usage:
+
+        %matplotlib widget
+        mask = DrawMaskMatplotlib(frame)
+        mask.ui()  # or mask.ui(vmin=..., vmax=...) to override auto-scaling
+
+        Left-click on the image to add vertices to the current region.
+        Click "Close region" to commit it (needs >= 3 vertices) and start a new one.
+        Click "Undo last vertex" / "Remove last region" to fix mistakes.
+
+        mask.finish()                          # commit whichever region is still open
+        mask.save(file_path)                   # write a mask description to a json file
+        scan_to_integrate.mask = mask.mask      # boolean array, True = masked out
+
+    The saved/loaded file format is the same one used by DrawMask.save()/.load() and by
+    PFGeneralIntegrator.loadPyHyperMask, so files are interchangeable between the two tools.
+    '''
+
+    def __init__(self, frame):
+        '''
+        Construct a DrawMaskMatplotlib object
+
+        Args:
+            frame (xarray or ndarray): a single data frame with pix_x and pix_y axes
+
+        '''
+        self.image_data = np.squeeze(frame.to_numpy() if hasattr(frame, "to_numpy") else np.asarray(frame))
+        if len(self.image_data.shape) > 2:
+            warnings.warn('This tool needs a single frame, not a stack!  .sel down to a single frame before starting!',stacklevel=2)
+        self.shape = self.image_data.shape  # (rows=pix_y, cols=pix_x)
+        self.polygons = []       # list of Nx2 (x, y) vertex arrays, already committed
+        self._current_verts = []
+        self._current_artist = None
+        self._fig = None
+        self._ax = None
+        self._cid = None
+        self._status = None
+
+    def ui(self, vmin=None, vmax=None):
+        '''
+        Draw the DrawMaskMatplotlib UI in a Jupyter notebook (requires %matplotlib widget).
+
+        Args:
+            vmin (float): optional lower bound for the log color scale (auto-scaled from the 1st percentile of positive finite pixels if omitted)
+            vmax (float): optional upper bound for the log color scale (auto-scaled from the 99th percentile of positive finite pixels if omitted)
+
+        '''
+        fig, ax = plt.subplots(figsize=(6, 6))
+        finite_positive = self.image_data[np.isfinite(self.image_data) & (self.image_data > 0)]
+        if finite_positive.size and (vmin is None or vmax is None):
+            auto_vmin, auto_vmax = np.percentile(finite_positive, [1, 99])
+            if vmin is None:
+                vmin = auto_vmin
+            if vmax is None:
+                vmax = auto_vmax
+        if finite_positive.size and vmin is not None and vmax is not None and vmax > vmin:
+            norm = LogNorm(vmin=vmin, vmax=vmax)
+        else:
+            norm = None  # no usable positive data range; fall back to linear auto-scaling
+        ax.imshow(self.image_data, norm=norm, origin="upper")
+        ax.set_title("Left-click to add vertices; 'Close region' to finish")
+
+        self._fig, self._ax = fig, ax
+        self._current_verts = []
+        (self._current_artist,) = ax.plot([], [], "o-", color="red", markersize=4)
+        self._cid = fig.canvas.mpl_connect("button_press_event", self._on_click)
+
+        close_button = widgets.Button(description="Close region")
+        undo_button = widgets.Button(description="Undo last vertex")
+        remove_button = widgets.Button(description="Remove last region")
+        self._status = widgets.Label(value=f"Regions saved: {len(self.polygons)}")
+        close_button.on_click(self._close_region)
+        undo_button.on_click(self._undo_vertex)
+        remove_button.on_click(self._remove_last)
+
+        display(widgets.HBox([close_button, undo_button, remove_button]))
+        display(self._status)
+        plt.show()
+
+    def _on_click(self, event):
+        if event.inaxes != self._ax or event.button != 1 or event.xdata is None:
+            return
+        self._current_verts.append((event.xdata, event.ydata))
+        self._redraw_current()
+
+    def _redraw_current(self):
+        xs, ys = zip(*self._current_verts) if self._current_verts else ([], [])
+        self._current_artist.set_data(xs, ys)
+        self._fig.canvas.draw_idle()
+
+    def _undo_vertex(self, _btn=None):
+        if self._current_verts:
+            self._current_verts.pop()
+            self._redraw_current()
+
+    def _close_region(self, _btn=None):
+        if len(self._current_verts) >= 3:
+            self.polygons.append(np.array(self._current_verts))
+            self._current_verts = []
+            self._redraw_current()
+            self._status.value = f"Regions saved: {len(self.polygons)}"
+
+    def _remove_last(self, _btn=None):
+        if self.polygons:
+            self.polygons.pop()
+            self._status.value = f"Regions saved: {len(self.polygons)}"
+
+    def finish(self):
+        '''
+        Commit whichever region is currently being drawn (call before save()).
+        '''
+        self._close_region()
+
+    @property
+    def mask(self):
+        '''
+        Render the mask as a numpy boolean array.
+        '''
+        ny, nx = self.shape
+        yy, xx = np.mgrid[0:ny, 0:nx]
+        points = np.column_stack((xx.ravel(), yy.ravel()))
+        mask_arr = np.zeros(self.shape, dtype=bool)
+        for verts in self.polygons:
+            inside = MplPath(verts).contains_points(points).reshape(ny, nx)
+            mask_arr |= inside
+        return mask_arr
+
+    def save(self, file_path):
+        '''
+        Save a parametric mask description as a json dump file, in the same format used by
+        DrawMask.save() and consumed by PFGeneralIntegrator.loadPyHyperMask.
+
+        Args:
+            file_path (str): name of the file to save
+
+        '''
+        records = []
+        for verts in self.polygons:
+            xs, ys = verts[:, 0], verts[:, 1]
+            records.append(json.dumps({
+                "x": {str(i): float(x) for i, x in enumerate(xs)},
+                "y": {str(i): float(y) for i, y in enumerate(ys)},
+            }))
+        with open(file_path, "w") as f:
+            json.dump(records, f)
+
+    def load(self, file_path):
+        '''
+        Load polygon vertices from an existing mask description file (reads DrawMask-format
+        files too).
+
+        Args:
+            file_path (str): name of the file to read from
+
+        '''
+        with open(file_path, "r") as f:
+            raw = json.load(f)
+        if isinstance(raw, dict):
+            raw = [json.dumps(raw)]
+        self.polygons = []
+        for item in raw:
+            record = json.loads(item) if isinstance(item, str) else item
+            xs = [record["x"][k] for k in sorted(record["x"], key=int)]
+            ys = [record["y"][k] for k in sorted(record["y"], key=int)]
+            self.polygons.append(np.column_stack((xs, ys)))
+
+
 class CMSGIWAXS:
     """For streamlined loading for CMS data"""
     def __init__(self, files, loader, integrator):

tests/test_DrawMaskMatplotlib.py

@@ -0,0 +1,62 @@
+import numpy as np
+import pytest
+
+pytest.importorskip("ipywidgets")
+
+from PyHyperScattering.IntegrationUtils import DrawMaskMatplotlib
+from PyHyperScattering.PFGeneralIntegrator import PFGeneralIntegrator
+
+
+def _square_polygon():
+    # a 5x5 square from (2,2) to (7,7) in (x, y), matching the fixture used in
+    # test_PFGeneralIntegrator_mask_parsing.py
+    return np.array([[2, 2], [7, 2], [7, 7], [2, 7]])
+
+
+def test_mask_property_marks_polygon_interior():
+    drawer = DrawMaskMatplotlib(np.zeros((10, 10)))
+    drawer.polygons.append(_square_polygon())
+
+    mask = drawer.mask
+
+    assert mask.shape == (10, 10)
+    assert mask.dtype == bool
+    assert mask.any()
+    assert mask[4, 4]  # inside the square
+    assert not mask[0, 0]  # outside the square
+
+
+def test_save_and_load_round_trips_polygons(tmp_path):
+    drawer = DrawMaskMatplotlib(np.zeros((10, 10)))
+    drawer.polygons.append(_square_polygon())
+
+    mask_file = tmp_path / "mask.json"
+    drawer.save(mask_file)
+
+    reloaded = DrawMaskMatplotlib(np.zeros((10, 10)))
+    reloaded.load(mask_file)
+
+    assert len(reloaded.polygons) == 1
+    np.testing.assert_allclose(reloaded.polygons[0], drawer.polygons[0])
+    np.testing.assert_array_equal(reloaded.mask, drawer.mask)
+
+
+def test_saved_mask_file_loads_in_pfgeneralintegrator(tmp_path):
+    drawer = DrawMaskMatplotlib(np.zeros((10, 10)))
+    drawer.polygons.append(_square_polygon())
+
+    mask_file = tmp_path / "mask.json"
+    drawer.save(mask_file)
+
+    integrator = PFGeneralIntegrator(maskmethod="none", geomethod="none")
+    integrator.loadPyHyperMask(maskpath=mask_file, maskshape=(10, 10))
+
+    assert integrator.mask.shape == (10, 10)
+    assert integrator.mask.dtype == bool
+    # PFGeneralIntegrator.mask uses the opposite convention from DrawMaskMatplotlib.mask:
+    # True = valid/unmasked pixel, vs. True = masked out. Boundary pixels can differ by one
+    # pixel between skimage.draw.polygon2mask (used by PFGeneralIntegrator) and
+    # matplotlib.path.Path.contains_points (used by DrawMaskMatplotlib), so only compare
+    # points clearly inside/outside the polygon rather than exact edge-pixel agreement.
+    assert not integrator.mask[4, 4]  # inside the square: masked out -> not valid
+    assert integrator.mask[0, 0]  # outside the square: not masked -> valid