diff --git a/pyproject.toml b/pyproject.toml index 5650d267..11835da3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -55,7 +55,9 @@ ui = [ "holoviews", "hvplot", "ipyaggrid", - "matplotlib" + "matplotlib", + "ipympl", + "ipywidgets" ] doc = [ "sphinx", diff --git a/src/PyHyperScattering/IntegrationUtils.py b/src/PyHyperScattering/IntegrationUtils.py index 123981db..10d3f3dd 100644 --- a/src/PyHyperScattering/IntegrationUtils.py +++ b/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): diff --git a/tests/test_DrawMaskMatplotlib.py b/tests/test_DrawMaskMatplotlib.py new file mode 100644 index 00000000..2c5ce2ad --- /dev/null +++ b/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