diff --git a/probeflow/gui/dialogs/fft_selection.py b/probeflow/gui/dialogs/fft_selection.py index 70a7131..890ae98 100644 --- a/probeflow/gui/dialogs/fft_selection.py +++ b/probeflow/gui/dialogs/fft_selection.py @@ -1,11 +1,17 @@ -"""Interactive circle / ellipse Fourier selections on the FFT viewer axes. +"""Interactive Fourier selections on the FFT viewer axes. -Mirrors the drag/hit-test pattern of -:class:`probeflow.gui.lattice_grid.fft_overlay.FFTLatticeOverlay`: geometry is -held in q-space (nm⁻¹), drawn as matplotlib patches on the FFT axes, and -hit-tested in display pixels. Each selection's conjugate partner (point -reflection through DC) is drawn dashed — it is what makes the inverse-FFT -reconstruction real-valued. +Three selection kinds, all created by **drawing** on the FFT (drag from one +corner to the opposite, as ROIs are drawn elsewhere in the app): + +* ``ellipse`` — drag a bounding box; hold **Shift** for a circle. +* ``rect`` — drag a bounding box; hold **Shift** for a square. +* ``paint`` — freehand brush; the union of circular stamps forms the region. + +Geometry is held in q-space (nm⁻¹), drawn as matplotlib patches / an RGBA +``AxesImage`` on the FFT axes, and hit-tested in display pixels. Each +selection's conjugate partner (point reflection through DC) is shown — dashed +for shapes, mirrored for paint — it is what makes the inverse-FFT reconstruction +real-valued. The overlay is GUI-state only; mask construction and the inverse transform live in :mod:`probeflow.processing.inverse_fft`. @@ -14,26 +20,37 @@ from __future__ import annotations import math -from dataclasses import dataclass +from dataclasses import dataclass, field import numpy as np -from matplotlib.patches import Ellipse +from matplotlib.patches import Ellipse, Rectangle _HIT_RADIUS_PX = 10.0 _HANDLE_NONE = 0 _HANDLE_CENTRE = 1 _HANDLE_RX = 2 _HANDLE_RY = 3 +_MIN_DRAW_PX = 3.0 # drags smaller than this (an accidental click) are discarded + +_DEFAULT_PAINT_COLOR = (137, 220, 235) # #89dceb — matches the selection cyan @dataclass class FourierSelection: - """A circle/ellipse selection, centre + semi-axes in q-space (nm⁻¹).""" - kind: str # "circle" | "ellipse" - cx_q: float - cy_q: float - rx_q: float - ry_q: float + """A Fourier selection in q-space (nm⁻¹). + + ``rx_q``/``ry_q`` are the semi-axes for an ellipse or the half-extents for a + rectangle. Paint selections instead use ``stamps_q`` (circular-brush centres) + + ``radius_q`` (brush radius) + ``color`` (RGB 0–255). + """ + kind: str # "ellipse" | "rect" | "paint" + cx_q: float = 0.0 + cy_q: float = 0.0 + rx_q: float = 0.0 + ry_q: float = 0.0 + stamps_q: list = field(default_factory=list) + radius_q: float = 0.0 + color: tuple = _DEFAULT_PAINT_COLOR class FFTSelectionOverlay: @@ -50,6 +67,15 @@ def __init__(self, ax, qx, qy, image_shape, on_change=None): self._artists: list = [] self._drag_handle = _HANDLE_NONE self._drag_idx = -1 + self._drag_last = None # last cursor (x_q, y_q) for body/paint drags + # Draw-to-create state. + self._tool: str | None = None # None | "ellipse" | "rect" | "paint" + self._draw_anchor = None # (x_q, y_q) anchor for an in-progress shape + self._drawing = False + self._painting = False + # Paint controls. + self._brush_px = 8.0 + self._paint_color = _DEFAULT_PAINT_COLOR # ── geometry helpers ────────────────────────────────────────────────────── @@ -64,19 +90,28 @@ def set_qaxes(self, qx, qy, image_shape) -> None: self._qy = np.asarray(qy, dtype=np.float64) self._shape = (int(image_shape[0]), int(image_shape[1])) - def _default_radius_q(self) -> float: - # ~6 % of the qx half-range, a comfortable starting spot size. - half = max(abs(float(self._qx[0])), abs(float(self._qx[-1]))) - return max(half * 0.06, 3.0 * abs(self._dq()[0])) + def _brush_radius_q(self) -> float: + return self._brush_px * abs(self._dq()[0]) - # ── public API ──────────────────────────────────────────────────────────── + # ── tool / paint configuration ──────────────────────────────────────────── - def add(self, kind: str = "circle") -> None: - r = self._default_radius_q() - # Place a little off-DC so it is visible and not on the central spot. - self._sels.append(FourierSelection(kind, cx_q=3 * r, cy_q=0.0, rx_q=r, ry_q=r)) - self._selected = len(self._sels) - 1 - self._notify() + def set_tool(self, tool: str | None) -> None: + self._tool = tool if tool in ("ellipse", "rect", "paint") else None + + def tool(self) -> str | None: + return self._tool + + def set_brush_radius_px(self, px: float) -> None: + self._brush_px = max(0.5, float(px)) + + def set_paint_color(self, rgb) -> None: + self._paint_color = tuple(int(c) for c in rgb) + if (0 <= self._selected < len(self._sels) + and self._sels[self._selected].kind == "paint"): + self._sels[self._selected].color = self._paint_color + self._notify() + + # ── public API ──────────────────────────────────────────────────────────── def delete_selected(self) -> None: if 0 <= self._selected < len(self._sels): @@ -93,21 +128,39 @@ def clear(self) -> None: def count(self) -> int: return len(self._sels) - def to_fft_ellipses(self) -> list[dict]: - """Selections as FFT-pixel ellipse dicts {dx,dy,rx,ry,angle_deg} (+ q for - provenance), for mask building / the inverse_fft_filter op.""" + def to_regions(self) -> list[dict]: + """Selections as FFT-pixel region dicts (+ q-space provenance), for mask + building and the ``inverse_fft_filter`` op. Each dict carries a ``kind`` + of ``"ellipse"``, ``"rect"`` or ``"paint"``.""" dqx, dqy = self._dq() out: list[dict] = [] for s in self._sels: - out.append({ + if s.kind == "paint": + stamps = [[x / dqx if dqx else 0.0, y / dqy if dqy else 0.0] + for (x, y) in s.stamps_q] + out.append({ + "kind": "paint", + "stamps": stamps, + "radius": max(s.radius_q / abs(dqx), 0.5) if dqx else 1.0, + "stamps_q": [list(p) for p in s.stamps_q], + "radius_q": s.radius_q, + "color": list(s.color), + }) + continue + d = { + "kind": s.kind, "dx": s.cx_q / dqx if dqx else 0.0, "dy": s.cy_q / dqy if dqy else 0.0, - "rx": max(s.rx_q / abs(dqx), 0.5) if dqx else 1.0, - "ry": max(s.ry_q / abs(dqy), 0.5) if dqy else 1.0, "angle_deg": 0.0, "cx_q": s.cx_q, "cy_q": s.cy_q, "rx_q": s.rx_q, "ry_q": s.ry_q, - "kind": s.kind, - }) + } + if s.kind == "rect": + d["half_w"] = max(s.rx_q / abs(dqx), 0.5) if dqx else 1.0 + d["half_h"] = max(s.ry_q / abs(dqy), 0.5) if dqy else 1.0 + else: # ellipse + d["rx"] = max(s.rx_q / abs(dqx), 0.5) if dqx else 1.0 + d["ry"] = max(s.ry_q / abs(dqy), 0.5) if dqy else 1.0 + out.append(d) return out # ── drawing ──────────────────────────────────────────────────────────────── @@ -117,13 +170,22 @@ def draw(self) -> None: self._artists = [] # axes were cla()'d by the host redraw for i, s in enumerate(self._sels): chosen = (i == self._selected) + if s.kind == "paint": + self._draw_paint(s, chosen) + continue + edge = "#89dceb" if chosen else "#89b4fa" + lw = 1.6 if chosen else 1.2 for cx, cy, dashed in ((s.cx_q, s.cy_q, False), (-s.cx_q, -s.cy_q, True)): - e = Ellipse((cx, cy), 2 * s.rx_q, 2 * s.ry_q, angle=0.0, - fill=False, edgecolor="#89dceb" if chosen else "#89b4fa", - lw=1.6 if chosen else 1.2, - ls="--" if dashed else "-", zorder=11) - self._ax.add_patch(e) - self._artists.append(e) + if s.kind == "rect": + patch = Rectangle((cx - s.rx_q, cy - s.ry_q), 2 * s.rx_q, 2 * s.ry_q, + fill=False, edgecolor=edge, lw=lw, + ls="--" if dashed else "-", zorder=11) + else: + patch = Ellipse((cx, cy), 2 * s.rx_q, 2 * s.ry_q, angle=0.0, + fill=False, edgecolor=edge, lw=lw, + ls="--" if dashed else "-", zorder=11) + self._ax.add_patch(patch) + self._artists.append(patch) if chosen: for hx, hy in ((s.cx_q, s.cy_q), (s.cx_q + s.rx_q, s.cy_q), (s.cx_q, s.cy_q + s.ry_q)): @@ -131,15 +193,64 @@ def draw(self) -> None: markersize=7, markeredgewidth=1.0, zorder=12) self._artists.append(art) + def _draw_paint(self, s: FourierSelection, chosen: bool) -> None: + rgba = self._rasterize_paint(s, chosen) + if rgba is None: + return + extent = [float(self._qx[0]), float(self._qx[-1]), + float(self._qy[-1]), float(self._qy[0])] + # imshow can autoscale the axes; keep the current view (zoom/pan) fixed. + xlim, ylim = self._ax.get_xlim(), self._ax.get_ylim() + im = self._ax.imshow(rgba, origin="upper", extent=extent, aspect="auto", + interpolation="nearest", zorder=10.5 if chosen else 10) + self._ax.set_xlim(*xlim) + self._ax.set_ylim(*ylim) + self._artists.append(im) + + def _rasterize_paint(self, s: FourierSelection, chosen: bool): + """Stamp the painted discs (and their conjugates) into an RGBA overlay + matching the FFT pixel grid, so it aligns with the applied mask.""" + if not s.stamps_q: + return None + ny, nx = self._shape + dqx, dqy = self._dq() + if not dqx or not dqy: + return None + cx0, cy0 = nx // 2, ny // 2 + rpx = max(s.radius_q / abs(dqx), 0.5) + rpy = max(s.radius_q / abs(dqy), 0.5) + mask = np.zeros((ny, nx), dtype=bool) + for (x, y) in s.stamps_q: + dx, dy = x / dqx, y / dqy + for sx, sy in ((dx, dy), (-dx, -dy)): # show the conjugate too + ix, iy = int(round(cx0 + sx)), int(round(cy0 + sy)) + x0, x1 = max(0, ix - int(rpx) - 1), min(nx, ix + int(rpx) + 2) + y0, y1 = max(0, iy - int(rpy) - 1), min(ny, iy + int(rpy) + 2) + if x0 >= x1 or y0 >= y1: + continue + ys, xs = np.ogrid[y0:y1, x0:x1] + disc = ((xs - ix) / rpx) ** 2 + ((ys - iy) / rpy) ** 2 <= 1.0 + mask[y0:y1, x0:x1] |= disc + if not mask.any(): + return None + rgba = np.zeros((ny, nx, 4), dtype=np.uint8) + rgba[mask] = (*[int(c) for c in s.color], 200 if chosen else 130) + return rgba + # ── hit-testing + drag ────────────────────────────────────────────────────── def _handle_at(self, event) -> tuple[int, int]: """Return (selection_index, handle_id) under the cursor, or (-1, NONE).""" - # Prefer handles of the currently-selected item, then any body. order = ([self._selected] if 0 <= self._selected < len(self._sels) else []) + \ [i for i in range(len(self._sels)) if i != self._selected] for i in order: s = self._sels[i] + if s.kind == "paint": + if (event.xdata is not None and event.ydata is not None + and any(math.hypot(event.xdata - x, event.ydata - y) <= s.radius_q + for (x, y) in s.stamps_q)): + return i, _HANDLE_CENTRE + continue handles = { _HANDLE_RX: (s.cx_q + s.rx_q, s.cy_q), _HANDLE_RY: (s.cx_q, s.cy_q + s.ry_q), @@ -152,29 +263,65 @@ def _handle_at(self, event) -> tuple[int, int]: return i, hid except Exception: pass - # Body hit (inside the ellipse, in q-space). + # Body hit (inside the shape, in q-space). if event.xdata is not None and event.ydata is not None: - dx = (event.xdata - s.cx_q) / max(s.rx_q, 1e-9) - dy = (event.ydata - s.cy_q) / max(s.ry_q, 1e-9) - if dx * dx + dy * dy <= 1.0: - return i, _HANDLE_CENTRE + if s.kind == "rect": + if (abs(event.xdata - s.cx_q) <= s.rx_q + and abs(event.ydata - s.cy_q) <= s.ry_q): + return i, _HANDLE_CENTRE + else: + dx = (event.xdata - s.cx_q) / max(s.rx_q, 1e-9) + dy = (event.ydata - s.cy_q) / max(s.ry_q, 1e-9) + if dx * dx + dy * dy <= 1.0: + return i, _HANDLE_CENTRE return -1, _HANDLE_NONE + def _shift_held(self, event) -> bool: + ge = getattr(event, "guiEvent", None) + if ge is not None: + try: + from PySide6.QtCore import Qt + return bool(ge.modifiers() & Qt.ShiftModifier) + except Exception: + pass + return "shift" in (getattr(event, "key", None) or "") + def on_press(self, event) -> bool: - """Begin a drag if a selection/handle is hit. Returns True if consumed.""" + """Begin editing a hit selection, or draw a new one if a tool is active. + Returns True if the event was consumed.""" if event.inaxes is not self._ax: return False idx, hid = self._handle_at(event) - if idx < 0: - return False - self._selected = idx - self._drag_idx = idx - self._drag_handle = hid + if idx >= 0: + self._selected = idx + self._drag_idx = idx + self._drag_handle = hid + self._drag_last = (event.xdata, event.ydata) + self._notify() + return True + if self._tool is not None and event.xdata is not None and event.ydata is not None: + return self._begin_draw(float(event.xdata), float(event.ydata)) + return False + + def _begin_draw(self, x: float, y: float) -> bool: + if self._tool == "paint": + sel = FourierSelection("paint", radius_q=self._brush_radius_q(), + color=self._paint_color) + sel.stamps_q.append((x, y)) + self._sels.append(sel) + self._selected = self._drag_idx = len(self._sels) - 1 + self._painting = True + else: + sel = FourierSelection(self._tool, cx_q=x, cy_q=y) + self._sels.append(sel) + self._selected = self._drag_idx = len(self._sels) - 1 + self._draw_anchor = (x, y) + self._drawing = True self._notify() return True def on_motion(self, event) -> bool: - if self._drag_handle == _HANDLE_NONE or self._drag_idx < 0: + if not (self._drawing or self._painting or self._drag_handle != _HANDLE_NONE): return False xdata, ydata = event.xdata, event.ydata if xdata is None or ydata is None: @@ -183,26 +330,82 @@ def on_motion(self, event) -> bool: (float(event.x), float(event.y))) except Exception: return True + xdata, ydata = float(xdata), float(ydata) + + if self._painting and 0 <= self._drag_idx < len(self._sels): + sel = self._sels[self._drag_idx] + step = max(sel.radius_q * 0.4, abs(self._dq()[0])) + if (not sel.stamps_q + or math.hypot(xdata - sel.stamps_q[-1][0], + ydata - sel.stamps_q[-1][1]) >= step): + sel.stamps_q.append((xdata, ydata)) + self._notify() + return True + + if self._drawing and 0 <= self._drag_idx < len(self._sels): + self._update_draw(self._sels[self._drag_idx], xdata, ydata, + self._shift_held(event)) + self._notify() + return True + s = self._sels[self._drag_idx] - if self._drag_handle == _HANDLE_CENTRE: - s.cx_q, s.cy_q = float(xdata), float(ydata) + if s.kind == "paint": # body drag: translate every stamp + if self._drag_last is not None: + lx, ly = self._drag_last + s.stamps_q = [(px + (xdata - lx), py + (ydata - ly)) + for (px, py) in s.stamps_q] + self._drag_last = (xdata, ydata) + elif self._drag_handle == _HANDLE_CENTRE: + s.cx_q, s.cy_q = xdata, ydata elif self._drag_handle == _HANDLE_RX: - s.rx_q = max(abs(float(xdata) - s.cx_q), abs(self._dq()[0])) - if s.kind == "circle": + s.rx_q = max(abs(xdata - s.cx_q), abs(self._dq()[0])) + if self._shift_held(event): s.ry_q = s.rx_q elif self._drag_handle == _HANDLE_RY: - s.ry_q = max(abs(float(ydata) - s.cy_q), abs(self._dq()[1])) - if s.kind == "circle": + s.ry_q = max(abs(ydata - s.cy_q), abs(self._dq()[1])) + if self._shift_held(event): s.rx_q = s.ry_q self._notify() return True + def _update_draw(self, s: FourierSelection, x: float, y: float, regular: bool) -> None: + ax0, ay0 = self._draw_anchor + sx, sy = x - ax0, y - ay0 + if regular: # Shift → circle / square + m = max(abs(sx), abs(sy)) + sx = math.copysign(m, sx) if sx else m + sy = math.copysign(m, sy) if sy else m + s.cx_q = ax0 + sx / 2.0 + s.cy_q = ay0 + sy / 2.0 + # No size floor here: a zero drag stays zero so on_release discards an + # accidental click. to_regions()/the mask builder floor the final size. + s.rx_q = abs(sx) / 2.0 + s.ry_q = abs(sy) / 2.0 + def on_release(self, event) -> None: + finishing_draw = self._drawing + idx = self._drag_idx self._drag_handle = _HANDLE_NONE self._drag_idx = -1 + self._drag_last = None + self._draw_anchor = None + self._drawing = False + self._painting = False + # Discard a shape that was only clicked, not dragged. + if finishing_draw and 0 <= idx < len(self._sels): + s = self._sels[idx] + try: + p0 = self._ax.transData.transform((s.cx_q - s.rx_q, s.cy_q - s.ry_q)) + p1 = self._ax.transData.transform((s.cx_q + s.rx_q, s.cy_q + s.ry_q)) + if math.hypot(p1[0] - p0[0], p1[1] - p0[1]) < _MIN_DRAW_PX: + self._sels.pop(idx) + self._selected = min(self._selected, len(self._sels) - 1) + self._notify() + except Exception: + pass def is_dragging(self) -> bool: - return self._drag_handle != _HANDLE_NONE + return (self._drag_handle != _HANDLE_NONE or self._drawing or self._painting) def _notify(self) -> None: if self._on_change is not None: diff --git a/probeflow/gui/dialogs/fft_viewer_reconstruct_mixin.py b/probeflow/gui/dialogs/fft_viewer_reconstruct_mixin.py index 4e473fa..0d14f9d 100644 --- a/probeflow/gui/dialogs/fft_viewer_reconstruct_mixin.py +++ b/probeflow/gui/dialogs/fft_viewer_reconstruct_mixin.py @@ -12,20 +12,33 @@ from __future__ import annotations import numpy as np +from probeflow.core.resources import asset_path from probeflow.gui._tooltips import tip as _tip from probeflow.gui.typography import ui_font +from PySide6.QtCore import Qt, QSize +from PySide6.QtGui import QIcon from PySide6.QtWidgets import ( QCheckBox, QComboBox, QFileDialog, QFrame, QGroupBox, QHBoxLayout, - QLabel, QPushButton, QScrollArea, QSpinBox, QVBoxLayout, QWidget, + QLabel, QPushButton, QScrollArea, QSpinBox, QToolButton, QVBoxLayout, QWidget, ) +# Shared paint palette (mirrors the feature-finder MASK_COLORS). +_PAINT_COLORS: dict[str, tuple[int, int, int]] = { + "Cyan": (137, 220, 235), + "Red": (243, 139, 168), + "Green": (166, 227, 161), + "Yellow": (249, 226, 175), + "Blue": (137, 180, 250), + "Magenta": (203, 166, 247), +} + class FFTViewerReconstructMixin: """Select Fourier features, preview the inverse FFT, and apply/export.""" def _build_reconstruct_tab(self) -> QWidget: - """Select Fourier features (circle/ellipse), preview the inverse FFT - result + residual, and apply/export.""" + """Draw Fourier selections (ellipse / rectangle / paint), preview the + inverse FFT result + residual, and apply/export.""" scroll = QScrollArea() scroll.setWidgetResizable(True) scroll.setFrameShape(QFrame.NoFrame) @@ -41,9 +54,9 @@ def _build_reconstruct_tab(self) -> QWidget: intro.setWordWrap(True) intro.setFont(ui_font(9)) intro.setToolTip(_tip( - "Make FFT filtering auditable: drop circle/ellipse selections on " - "Fourier features, choose Remove or Keep, and preview the " - "reconstructed image and the residual before applying. Use it to " + "Make FFT filtering auditable: draw ellipse/rectangle/paint " + "selections on Fourier features, choose Remove or Keep, and preview " + "the reconstructed image and the residual before applying. Use it to " "confirm a periodic artefact is really gone (or to isolate one " "periodic component).")) lay.addWidget(intro) @@ -52,31 +65,70 @@ def _build_reconstruct_tab(self) -> QWidget: sgrp = QGroupBox("Fourier selections") sg = QVBoxLayout(sgrp) sg.setSpacing(4) - add_row = QHBoxLayout() - add_circle = QPushButton("Add circle") - add_circle.setToolTip(_tip( - "Add a circular selection. Drag its centre to move it onto a " - "Fourier feature and drag the square handle to resize. Its " - "conjugate partner (dashed) is added automatically.")) - add_circle.clicked.connect(lambda: self._on_add_selection("circle")) - add_ellipse = QPushButton("Add ellipse") - add_ellipse.setToolTip(_tip( - "Add an elliptical selection with independent width/height handles " - "— for elongated or streaky Fourier features.")) - add_ellipse.clicked.connect(lambda: self._on_add_selection("ellipse")) + # Checkable draw tools: pick one, then drag on the FFT to draw. Holding + # Shift draws a regular shape (circle / square). Click an active tool + # again to return to edit mode (move/resize existing selections). + tool_row = QHBoxLayout() + self._recon_tool_btns: dict[str, QToolButton] = {} + tool_specs = [ + ("ellipse", "Ellipse", "ellipse", + "Draw an ellipse: drag a box on the FFT. Hold Shift for a circle. " + "Its conjugate partner (dashed) is added automatically."), + ("rect", "Rectangle", "rectangle", + "Draw a rectangle: drag a box on the FFT. Hold Shift for a square."), + ("paint", "Paint", "freehand", + "Freehand brush: drag to paint an irregular Fourier region. The " + "mirrored (conjugate) region is grabbed too."), + ] + for kind, label, icon, tip in tool_specs: + btn = QToolButton() + btn.setText(label) + btn.setCheckable(True) + btn.setToolButtonStyle(Qt.ToolButtonTextBesideIcon) + ipath = asset_path(f"toolbar/{icon}.png") + if ipath.exists(): + btn.setIcon(QIcon(str(ipath))) + btn.setIconSize(QSize(16, 16)) + btn.setToolTip(_tip(tip)) + btn.clicked.connect(lambda _checked=False, k=kind: self._on_tool_clicked(k)) + self._recon_tool_btns[kind] = btn + tool_row.addWidget(btn) + tool_row.addStretch(1) + sg.addLayout(tool_row) + + # Paint brush controls — only relevant while the Paint tool is active. + self._recon_paint_row = QWidget() + paint_lay = QHBoxLayout(self._recon_paint_row) + paint_lay.setContentsMargins(0, 0, 0, 0) + paint_lay.setSpacing(4) + paint_lay.addWidget(QLabel("Brush:")) + self._recon_brush_spin = QSpinBox() + self._recon_brush_spin.setRange(1, 100) + self._recon_brush_spin.setValue(8) + self._recon_brush_spin.setSuffix(" px") + self._recon_brush_spin.setToolTip(_tip("Paint brush radius, in FFT pixels.")) + self._recon_brush_spin.setMaximumWidth(80) + self._recon_brush_spin.valueChanged.connect(self._on_brush_size_changed) + paint_lay.addWidget(self._recon_brush_spin) + paint_lay.addWidget(QLabel("Color:")) + self._recon_color_combo = QComboBox() + self._recon_color_combo.addItems(list(_PAINT_COLORS.keys())) + self._recon_color_combo.setToolTip(_tip("Paint overlay color.")) + self._recon_color_combo.setMaximumWidth(110) + self._recon_color_combo.currentIndexChanged.connect(self._on_paint_color_changed) + paint_lay.addWidget(self._recon_color_combo) + paint_lay.addStretch(1) + self._recon_paint_row.setVisible(False) + sg.addWidget(self._recon_paint_row) + del_btn = QPushButton("Delete selected") del_btn.setToolTip(_tip("Remove the currently-selected Fourier region.")) del_btn.clicked.connect(self._on_delete_selection) clr_btn = QPushButton("Clear selections") clr_btn.setToolTip(_tip("Remove all Fourier selections.")) clr_btn.clicked.connect(self._on_clear_selections) - # 2×2 button block at a compact width instead of full-width rows. - for b in (add_circle, add_ellipse, del_btn, clr_btn): + for b in (del_btn, clr_btn): b.setMaximumWidth(150) - add_row.addWidget(add_circle) - add_row.addWidget(add_ellipse) - add_row.addStretch(1) - sg.addLayout(add_row) del_row = QHBoxLayout() del_row.addWidget(del_btn) del_row.addWidget(clr_btn) @@ -228,13 +280,33 @@ def _on_selection_changed(self) -> None: self._canvas_fft.draw_idle() self._update_reconstruct_status() - def _on_add_selection(self, kind: str) -> None: + def _on_tool_clicked(self, kind: str) -> None: ov = self._ensure_selection_overlay() if ov is None: + self._recon_tool_btns[kind].setChecked(False) self._recon_status_lbl.setText("Load a scan first.") return - ov.add(kind) - self._on_selection_changed() + # Toggle: a second click on the active tool returns to edit mode. + active = None if ov.tool() == kind else kind + ov.set_tool(active) + for k, btn in self._recon_tool_btns.items(): + btn.setChecked(k == active) + self._recon_paint_row.setVisible(active == "paint") + if active == "paint": + ov.set_brush_radius_px(self._recon_brush_spin.value()) + ov.set_paint_color(_PAINT_COLORS[self._recon_color_combo.currentText()]) + self._update_reconstruct_status() + + def _on_brush_size_changed(self, value: int) -> None: + ov = self._fft_selection_overlay + if ov is not None: + ov.set_brush_radius_px(int(value)) + + def _on_paint_color_changed(self, _idx: int) -> None: + ov = self._fft_selection_overlay + if ov is not None: + ov.set_paint_color(_PAINT_COLORS[self._recon_color_combo.currentText()]) + self._on_selection_changed() def _on_delete_selection(self) -> None: if self._fft_selection_overlay is not None: @@ -254,11 +326,10 @@ def _compute_reconstruction(self, array): if ov is None or ov.count() == 0: return None from probeflow.processing.inverse_fft import ( - FourierEllipse, fourier_ellipse_mask, inverse_fft_from_mask) - sels = [FourierEllipse(dx=s["dx"], dy=s["dy"], rx=s["rx"], ry=s["ry"]) - for s in ov.to_fft_ellipses()] - mask = fourier_ellipse_mask( - array.shape, sels, + fourier_region_from_dict, fourier_region_mask, inverse_fft_from_mask) + regions = [fourier_region_from_dict(d) for d in ov.to_regions()] + mask = fourier_region_mask( + array.shape, regions, conjugate=self._recon_conj_cb.isChecked(), soft_px=float(self._recon_soft_spin.value())) return inverse_fft_from_mask(array, mask, mode=self._reconstruct_mode()) @@ -284,14 +355,11 @@ def _on_reconstruct_clear(self) -> None: def _reconstruct_op_params(self) -> dict: ov = self._fft_selection_overlay - sels = ov.to_fft_ellipses() if ov is not None else [] + # to_regions() already returns JSON-safe dicts carrying the per-kind + # geometry (ellipse/rect in FFT px, paint as a pixel stamp list) plus + # q-space provenance. params = { - "selections": [ - {"dx": s["dx"], "dy": s["dy"], "rx": s["rx"], "ry": s["ry"], - "angle_deg": 0.0, "cx_q": s["cx_q"], "cy_q": s["cy_q"], - "rx_q": s["rx_q"], "ry_q": s["ry_q"], "kind": s["kind"]} - for s in sels - ], + "selections": ov.to_regions() if ov is not None else [], "mode": self._reconstruct_mode(), "conjugate_symmetric": bool(self._recon_conj_cb.isChecked()), "soft_px": float(self._recon_soft_spin.value()), @@ -364,8 +432,10 @@ def _update_reconstruct_status(self, res=None) -> None: mode = "remove selected" if self._reconstruct_mode() == "remove_selected" else "keep selected" src = "ROI" if self._fft_source == "active_roi" else "whole image" if n == 0: - self._recon_status_lbl.setText( - f"FFT source: {src}. Add a circle or ellipse to begin.") + tool = ov.tool() if ov is not None else None + hint = ("Drag on the FFT to draw (Shift = regular shape)." + if tool else "Pick Ellipse, Rectangle or Paint to begin.") + self._recon_status_lbl.setText(f"FFT source: {src}. {hint}") return conj = " + conjugates" if self._recon_conj_cb.isChecked() else "" txt = (f"FFT mask: {n} region{'s' if n != 1 else ''}{conj} · " diff --git a/probeflow/processing/inverse_fft.py b/probeflow/processing/inverse_fft.py index ac9f8d3..9914f57 100644 --- a/probeflow/processing/inverse_fft.py +++ b/probeflow/processing/inverse_fft.py @@ -31,8 +31,12 @@ __all__ = [ "FourierEllipse", + "FourierRect", + "FourierStrokes", "InverseFFTResult", "fourier_ellipse_mask", + "fourier_region_mask", + "fourier_region_from_dict", "inverse_fft_from_mask", "inverse_fft_filter", ] @@ -53,6 +57,33 @@ class FourierEllipse: angle_deg: float = 0.0 +@dataclass(frozen=True) +class FourierRect: + """A rectangular Fourier selection, in FFT-pixel offsets from the centred DC. + + ``dx``/``dy`` are the centre offsets from DC; ``half_w``/``half_h`` are the + half-extents in pixels; ``angle_deg`` rotates the rectangle. A square is + ``half_w == half_h`` with ``angle_deg == 0``. + """ + dx: float + dy: float + half_w: float + half_h: float + angle_deg: float = 0.0 + + +@dataclass(frozen=True) +class FourierStrokes: + """A freehand (painted) Fourier selection. + + ``stamps`` is a list of ``(dx, dy)`` circular-brush centres in FFT-pixel + offsets from DC; ``radius`` is the shared brush radius in pixels. The union + of the stamped discs forms the region. + """ + stamps: tuple + radius: float + + @dataclass class InverseFFTResult: result: np.ndarray # reconstructed real-space image @@ -62,50 +93,141 @@ class InverseFFTResult: imag_residual_norm: float # ‖Im(ifft)‖ / ‖Re(ifft)‖ — ~0 when symmetric -def fourier_ellipse_mask( +def _soft_ramp(signed_px: np.ndarray, soft_px: float) -> np.ndarray: + """Linear ramp from 1 (inside, ``signed_px <= 0``) to 0 over ``soft_px``. + + ``signed_px`` is the signed distance to the region boundary in pixels + (negative inside, positive outside). With ``soft_px == 0`` this is a hard + edge; otherwise the boundary feathers over ``soft_px`` pixels (reduces + ringing in the reconstruction). + """ + if soft_px > 0: + return np.clip(1.0 - signed_px / float(soft_px), 0.0, 1.0) + return (signed_px <= 0.0).astype(np.float64) + + +def _ellipse_contribution(xx, yy, cx, cy, e, soft_px): + rx = max(float(e.rx), 1e-6) + ry = max(float(e.ry), 1e-6) + theta = math.radians(float(e.angle_deg)) + ct, st = math.cos(theta), math.sin(theta) + px, py = cx + e.dx, cy + e.dy + xr = (xx - px) * ct + (yy - py) * st + yr = -(xx - px) * st + (yy - py) * ct + d = np.sqrt((xr / rx) ** 2 + (yr / ry) ** 2) # 1.0 on the boundary + # Convert the normalised radius to an approximate signed distance in pixels. + signed_px = (d - 1.0) * max(rx, ry) + return _soft_ramp(signed_px, soft_px) + + +def _rect_contribution(xx, yy, cx, cy, r, soft_px): + hw = max(float(r.half_w), 1e-6) + hh = max(float(r.half_h), 1e-6) + theta = math.radians(float(r.angle_deg)) + ct, st = math.cos(theta), math.sin(theta) + px, py = cx + r.dx, cy + r.dy + xr = (xx - px) * ct + (yy - py) * st + yr = -(xx - px) * st + (yy - py) * ct + # Signed distance to the rectangle boundary (approximate, exterior-exact). + ox = np.abs(xr) - hw + oy = np.abs(yr) - hh + outside = np.sqrt(np.maximum(ox, 0.0) ** 2 + np.maximum(oy, 0.0) ** 2) + inside = np.minimum(np.maximum(ox, oy), 0.0) + signed_px = outside + inside + return _soft_ramp(signed_px, soft_px) + + +def _strokes_contribution(xx, yy, cx, cy, s, sign, soft_px): + radius = max(float(s.radius), 1e-6) + stamps = s.stamps if sign > 0 else [(-dx, -dy) for (dx, dy) in s.stamps] + if not stamps: + return np.zeros(xx.shape, dtype=np.float64) + # Nearest-stamp distance → signed distance to the painted region boundary. + nearest = np.full(xx.shape, np.inf, dtype=np.float64) + for dx, dy in stamps: + px, py = cx + float(dx), cy + float(dy) + nearest = np.minimum(nearest, np.sqrt((xx - px) ** 2 + (yy - py) ** 2)) + signed_px = nearest - radius + return _soft_ramp(signed_px, soft_px) + + +def fourier_region_mask( shape: tuple[int, int], - ellipses, + regions, *, conjugate: bool = True, soft_px: float = 0.0, ) -> np.ndarray: """Build a float selection mask (0–1) in fftshift-centred FFT space. - Each :class:`FourierEllipse` marks pixels inside it as 1. When - ``conjugate`` is True (the default, required for a real reconstruction) an - exact point-reflected partner is added for every ellipse. ``soft_px > 0`` - gives the boundary a linear cosine-like ramp of that pixel width (reduces - ringing); ``soft_px == 0`` gives a hard edge. + ``regions`` is a heterogeneous list of :class:`FourierEllipse`, + :class:`FourierRect`, and :class:`FourierStrokes`. Each marks the pixels it + covers as 1, accumulated with ``np.maximum``. When ``conjugate`` is True + (the default, required for a real reconstruction) an exact point-reflected + partner is added for every region. ``soft_px > 0`` feathers boundaries over + that pixel width (reduces ringing); ``soft_px == 0`` gives a hard edge. """ Ny, Nx = int(shape[0]), int(shape[1]) cy, cx = Ny // 2, Nx // 2 yy, xx = np.mgrid[:Ny, :Nx].astype(np.float64) mask = np.zeros((Ny, Nx), dtype=np.float64) - for e in ellipses: - centres = [(e.dx, e.dy)] + for region in regions: + if isinstance(region, FourierStrokes): + signs = (1, -1) if conjugate else (1,) + for sign in signs: + mask = np.maximum( + mask, _strokes_contribution(xx, yy, cx, cy, region, sign, soft_px)) + continue + # Ellipse / rect: build the primary, then mirror it through DC. + from dataclasses import replace + variants = [region] if conjugate: - centres.append((-e.dx, -e.dy)) # exact point reflection through DC - rx = max(float(e.rx), 1e-6) - ry = max(float(e.ry), 1e-6) - theta = math.radians(float(e.angle_deg)) - ct, st = math.cos(theta), math.sin(theta) - for sx, sy in centres: - px, py = cx + sx, cy + sy - xr = (xx - px) * ct + (yy - py) * st - yr = -(xx - px) * st + (yy - py) * ct - d = np.sqrt((xr / rx) ** 2 + (yr / ry) ** 2) # 1.0 on the boundary - if soft_px > 0: - # Linear ramp from 1 (inside) to 0 over ``soft_px`` past the edge. - ramp = 1.0 - (d - 1.0) / (float(soft_px) / max(rx, ry)) - contribution = np.clip(ramp, 0.0, 1.0) + variants.append(replace(region, dx=-region.dx, dy=-region.dy)) + for v in variants: + if isinstance(v, FourierRect): + contrib = _rect_contribution(xx, yy, cx, cy, v, soft_px) else: - contribution = (d <= 1.0).astype(np.float64) - mask = np.maximum(mask, contribution) + contrib = _ellipse_contribution(xx, yy, cx, cy, v, soft_px) + mask = np.maximum(mask, contrib) return mask +def fourier_ellipse_mask( + shape: tuple[int, int], + ellipses, + *, + conjugate: bool = True, + soft_px: float = 0.0, +) -> np.ndarray: + """Back-compat wrapper over :func:`fourier_region_mask` for ellipse lists.""" + return fourier_region_mask( + shape, ellipses, conjugate=conjugate, soft_px=soft_px) + + +def fourier_region_from_dict(d: dict): + """Rebuild a region object from a serialised dict (op params / overlay). + + The ``kind`` key selects the type; a missing ``kind`` defaults to + ``"ellipse"`` for backward compatibility with states saved before rect/paint + existed. + """ + kind = str(d.get("kind", "ellipse")) + if kind == "paint": + stamps = tuple((float(p[0]), float(p[1])) for p in d.get("stamps", [])) + return FourierStrokes(stamps=stamps, radius=float(d.get("radius", 1.0))) + if kind == "rect": + return FourierRect( + dx=float(d.get("dx", 0.0)), dy=float(d.get("dy", 0.0)), + half_w=float(d.get("half_w", 1.0)), half_h=float(d.get("half_h", 1.0)), + angle_deg=float(d.get("angle_deg", 0.0))) + return FourierEllipse( + dx=float(d.get("dx", 0.0)), dy=float(d.get("dy", 0.0)), + rx=float(d.get("rx", 1.0)), ry=float(d.get("ry", 1.0)), + angle_deg=float(d.get("angle_deg", 0.0))) + + def inverse_fft_from_mask( image: np.ndarray, mask: np.ndarray, @@ -156,21 +278,23 @@ def inverse_fft_from_mask( def inverse_fft_filter( image: np.ndarray, - ellipses, + regions, *, mode: str = "remove_selected", conjugate: bool = True, soft_px: float = 0.0, ) -> np.ndarray: - """Convenience: build the ellipse mask and return only the reconstructed image. + """Convenience: build the region mask and return only the reconstructed image. - This is the entry point the ``inverse_fft_filter`` ProcessingState op calls. - Returns the input unchanged when there are no selections. + ``regions`` is a list of :class:`FourierEllipse` / :class:`FourierRect` / + :class:`FourierStrokes`. This is the entry point the ``inverse_fft_filter`` + ProcessingState op calls. Returns the input unchanged when there are no + selections. """ a = np.asarray(image, dtype=np.float64) if a.ndim != 2: raise ValueError("inverse_fft_filter expects a 2-D image") - if not ellipses: + if not regions: return a.copy() - mask = fourier_ellipse_mask(a.shape, ellipses, conjugate=conjugate, soft_px=soft_px) + mask = fourier_region_mask(a.shape, regions, conjugate=conjugate, soft_px=soft_px) return inverse_fft_from_mask(a, mask, mode=mode).result diff --git a/probeflow/processing/state.py b/probeflow/processing/state.py index 81e6efb..2123317 100644 --- a/probeflow/processing/state.py +++ b/probeflow/processing/state.py @@ -565,18 +565,13 @@ def apply_processing_state( elif step.op == "inverse_fft_filter": # Selection geometry is stored in FFT-pixel offsets (exact for the # full image this op applies to); q-space values, if present, are - # provenance-only. Conjugate symmetry keeps the result real. - from probeflow.processing.inverse_fft import FourierEllipse - sels = [ - FourierEllipse( - dx=float(s.get("dx", 0.0)), dy=float(s.get("dy", 0.0)), - rx=float(s.get("rx", 1.0)), ry=float(s.get("ry", 1.0)), - angle_deg=float(s.get("angle_deg", 0.0)), - ) - for s in p.get("selections", []) - ] + # provenance-only. Each selection carries a ``kind`` (ellipse / rect + # / paint); a missing kind reads as ellipse for legacy states. + # Conjugate symmetry keeps the result real. + from probeflow.processing.inverse_fft import fourier_region_from_dict + regions = [fourier_region_from_dict(s) for s in p.get("selections", [])] a = _proc.inverse_fft_filter( - a, sels, + a, regions, mode=str(p.get("mode", "remove_selected")), conjugate=bool(p.get("conjugate_symmetric", True)), soft_px=float(p.get("soft_px", 0.0)), diff --git a/tests/test_fft_phase_gui.py b/tests/test_fft_phase_gui.py index 7b81155..1b2b034 100644 --- a/tests/test_fft_phase_gui.py +++ b/tests/test_fft_phase_gui.py @@ -139,7 +139,11 @@ def test_phase_view_does_not_alter_data_or_tools(self, qapp): assert np.array_equal(dlg._fft_mag, mag) # magnitude data untouched # Reconstruct still applies in phase view. dlg._tab_widget.setCurrentIndex(dlg._reconstruct_tab_index) - dlg._on_add_selection("circle") + from probeflow.gui.dialogs.fft_selection import FourierSelection + ov = dlg._ensure_selection_overlay() + ov._sels.append(FourierSelection("ellipse", cx_q=1.5, cy_q=0.0, rx_q=0.5, ry_q=0.5)) + ov._selected = 0 + dlg._on_selection_changed() dlg._on_reconstruct_apply() assert captured.get("op") == "inverse_fft_filter" dlg.deleteLater() diff --git a/tests/test_fft_selection_gui.py b/tests/test_fft_selection_gui.py index 03140bc..537451c 100644 --- a/tests/test_fft_selection_gui.py +++ b/tests/test_fft_selection_gui.py @@ -45,46 +45,122 @@ def _dialog(qapp, captured=None, roi=False, new_image_fn=None): return dlg, img +def _ev(ax, xq, yq, shift=False): + """A minimal stand-in for a matplotlib mouse event in the FFT axes.""" + disp = ax.transData.transform((xq, yq)) + + class _E: + inaxes = ax + xdata = xq + ydata = yq + x = float(disp[0]) + y = float(disp[1]) + button = 1 + guiEvent = None + key = "shift" if shift else None + return _E() + + +def _draw(ov, ax, kind, p0, p1, shift=False): + """Drag a shape from corner ``p0`` to ``p1`` (q-space).""" + ov.set_tool(kind) + ov.on_press(_ev(ax, *p0)) + ov.on_motion(_ev(ax, *p1, shift=shift)) + ov.on_release(_ev(ax, *p1, shift=shift)) + ov.set_tool(None) + + +def _paint(ov, ax, pts): + ov.set_tool("paint") + ov.on_press(_ev(ax, *pts[0])) + for p in pts[1:]: + ov.on_motion(_ev(ax, *p)) + ov.on_release(_ev(ax, *pts[-1])) + ov.set_tool(None) + + +def _seed(dlg, kind="ellipse"): + """Add one selection without driving the canvas — for reconstruct-tab tests.""" + from probeflow.gui.dialogs.fft_selection import FourierSelection + ov = dlg._ensure_selection_overlay() + ov._sels.append(FourierSelection(kind, cx_q=1.5, cy_q=0.0, rx_q=0.5, ry_q=0.5)) + ov._selected = len(ov._sels) - 1 + dlg._on_selection_changed() + return ov + + class TestSelectionOverlay: - def test_add_move_resize_delete(self, qapp): + def test_draw_move_resize_delete(self, qapp): dlg, _img = _dialog(qapp) - dlg._on_add_selection("circle") - ov = dlg._fft_selection_overlay + ov = dlg._ensure_selection_overlay() + _draw(ov, dlg._ax_fft, "ellipse", (1.0, 0.5), (2.0, 1.5)) assert ov.count() == 1 and len(ov._artists) > 0 - # to_fft_ellipses maps q→px; centre is off-DC, radius ≥ 1 px. - e = ov.to_fft_ellipses()[0] - assert e["dx"] > 0 and e["rx"] >= 1.0 and e["kind"] == "circle" + e = ov.to_regions()[0] + assert e["kind"] == "ellipse" and e["dx"] > 0 and e["rx"] >= 1.0 # resize: enlarging rx_q grows the px radius ov._sels[0].rx_q *= 2 - assert ov.to_fft_ellipses()[0]["rx"] > e["rx"] + assert ov.to_regions()[0]["rx"] > e["rx"] ov.delete_selected() assert ov.count() == 0 + def test_tiny_click_is_discarded(self, qapp): + dlg, _ = _dialog(qapp) + ov = dlg._ensure_selection_overlay() + _draw(ov, dlg._ax_fft, "ellipse", (1.0, 0.5), (1.0, 0.5)) + assert ov.count() == 0 # a click without a drag makes no selection + def test_ellipse_independent_axes(self, qapp): dlg, _ = _dialog(qapp) - dlg._on_add_selection("ellipse") - ov = dlg._fft_selection_overlay - ov._sels[0].rx_q = 0.5 - ov._sels[0].ry_q = 0.2 - e = ov.to_fft_ellipses()[0] - assert e["rx"] != e["ry"] + ov = dlg._ensure_selection_overlay() + _draw(ov, dlg._ax_fft, "ellipse", (1.0, 0.2), (3.0, 0.6)) + e = ov.to_regions()[0] + assert e["rx"] != e["ry"] # free drag → unequal semi-axes + + def test_shift_makes_regular(self, qapp): + dlg, _ = _dialog(qapp) + ov = dlg._ensure_selection_overlay() + _draw(ov, dlg._ax_fft, "ellipse", (1.0, 0.2), (3.0, 0.6), shift=True) + e = ov.to_regions()[0] + assert e["rx"] == pytest.approx(e["ry"]) # Shift → circle + + def test_rectangle_kind_and_patches(self, qapp): + from matplotlib.patches import Rectangle + dlg, _ = _dialog(qapp) + ov = dlg._ensure_selection_overlay() + _draw(ov, dlg._ax_fft, "rect", (-2.0, -1.0), (-1.0, -0.2)) + r = ov.to_regions()[0] + assert r["kind"] == "rect" and r["half_w"] >= 1.0 and r["half_h"] >= 1.0 + # rectangle patch + its conjugate + assert sum(isinstance(a, Rectangle) for a in ov._artists) == 2 + + def test_paint_region(self, qapp): + from matplotlib.image import AxesImage + dlg, _ = _dialog(qapp) + ov = dlg._ensure_selection_overlay() + ov.set_brush_radius_px(6) + _paint(ov, dlg._ax_fft, [(0.5, -2.0), (1.0, -1.8), (1.5, -1.5), (2.0, -1.0)]) + assert ov.count() == 1 + p = ov.to_regions()[0] + assert p["kind"] == "paint" and len(p["stamps"]) >= 1 and p["radius"] > 0 + # the painted overlay is drawn as an image + assert any(isinstance(a, AxesImage) for a in ov._artists) def test_conjugate_drawn(self, qapp): dlg, _ = _dialog(qapp) - dlg._on_add_selection("circle") - ov = dlg._fft_selection_overlay - # two ellipse patches per selection (feature + conjugate) + 3 handles + ov = dlg._ensure_selection_overlay() + _draw(ov, dlg._ax_fft, "ellipse", (1.0, 0.5), (2.0, 1.5)) + # two ellipse patches per selection (feature + conjugate) from matplotlib.patches import Ellipse - n_ell = sum(isinstance(a, Ellipse) for a in ov._artists) - assert n_ell == 2 + assert sum(isinstance(a, Ellipse) for a in ov._artists) == 2 def test_clear(self, qapp): dlg, _ = _dialog(qapp) - dlg._on_add_selection("circle") - dlg._on_add_selection("ellipse") - assert dlg._fft_selection_overlay.count() == 2 + ov = dlg._ensure_selection_overlay() + _draw(ov, dlg._ax_fft, "ellipse", (1.0, 0.5), (2.0, 1.5)) + _draw(ov, dlg._ax_fft, "rect", (-2.0, -1.0), (-1.0, -0.2)) + assert ov.count() == 2 dlg._on_clear_selections() - assert dlg._fft_selection_overlay.count() == 0 + assert ov.count() == 0 class TestReconstructTab: @@ -96,7 +172,7 @@ def test_tab_present(self, qapp): def test_preview_does_not_mutate_and_sets_active(self, qapp): dlg, img = _dialog(qapp) - dlg._on_add_selection("circle") + _seed(dlg) before = img.copy() dlg._on_reconstruct_preview() assert dlg._reconstruct_preview_active is True @@ -107,7 +183,7 @@ def test_preview_does_not_mutate_and_sets_active(self, qapp): def test_whole_image_apply_routes_op(self, qapp): captured: dict = {} dlg, _ = _dialog(qapp, captured=captured) - dlg._on_add_selection("circle") + _seed(dlg) dlg._recon_mode_combo.setCurrentIndex(1) # Keep selected dlg._recon_soft_spin.setValue(2) dlg._on_reconstruct_apply() @@ -132,7 +208,7 @@ def test_roi_apply_degrades_without_new_image_host(self, qapp): # switch FFT source to the ROI dlg._fft_source = "active_roi" dlg._arr, dlg._scan_range_m = dlg._resolve_source_array() - dlg._on_add_selection("circle") + _seed(dlg) dlg._on_reconstruct_apply() assert "op" not in captured # no whole-image op routed assert "export" in dlg._recon_status_lbl.text().lower() @@ -150,7 +226,7 @@ def _open_new_image(arr, scan_range_m, provenance): dlg._fft_source = "active_roi" dlg._arr, dlg._scan_range_m = dlg._resolve_source_array() dlg._recompute_fft() - dlg._on_add_selection("circle") + _seed(dlg) dlg._on_reconstruct_apply() assert "op" not in captured @@ -162,7 +238,7 @@ def _open_new_image(arr, scan_range_m, provenance): def test_export_writes_file(self, qapp, monkeypatch, tmp_path): dlg, _ = _dialog(qapp) - dlg._on_add_selection("circle") + _seed(dlg) out = tmp_path / "result.csv" monkeypatch.setattr( "PySide6.QtWidgets.QFileDialog.getSaveFileName", diff --git a/tests/test_inverse_fft.py b/tests/test_inverse_fft.py index 6d03348..78f256e 100644 --- a/tests/test_inverse_fft.py +++ b/tests/test_inverse_fft.py @@ -7,7 +7,11 @@ from probeflow.processing.inverse_fft import ( FourierEllipse, + FourierRect, + FourierStrokes, fourier_ellipse_mask, + fourier_region_from_dict, + fourier_region_mask, inverse_fft_filter, inverse_fft_from_mask, ) @@ -59,6 +63,68 @@ def test_ellipse_anisotropy(self): assert m[CY + 6, CX] == 0.0 # narrow along y +# ─── fourier_region_mask: rect / paint ─────────────────────────────────────── + +class TestRegionMask: + def test_ellipse_wrapper_matches_region(self): + m_old = fourier_ellipse_mask((N, N), _sine_circle()) + m_new = fourier_region_mask((N, N), _sine_circle()) + assert np.array_equal(m_old, m_new) + + def test_rect_covers_box_and_conjugate(self): + r = FourierRect(dx=KX, dy=0, half_w=4, half_h=2) + m = fourier_region_mask((N, N), [r]) + assert m[CY, CX + KX] == 1.0 + assert m[CY, CX - KX] == 1.0 # conjugate + assert m[CY, CX] == 0.0 # DC untouched + assert m[CY, CX + KX + 5] == 0.0 # outside the half-width (4) + + def test_square_is_symmetric(self): + sq = FourierRect(dx=0, dy=0, half_w=5, half_h=5) + m = fourier_region_mask((N, N), [sq]) + assert m[CY, CX + 5] == 1.0 and m[CY + 5, CX] == 1.0 + assert m[CY, CX + 7] == 0.0 and m[CY + 7, CX] == 0.0 + + def test_strokes_cover_stamps_and_conjugate(self): + s = FourierStrokes(stamps=((KX, 0), (KX + 2, 0)), radius=3) + m = fourier_region_mask((N, N), [s]) + assert m[CY, CX + KX] == 1.0 + assert m[CY, CX - KX] == 1.0 # conjugate stamp + assert m[CY + 20, CX] == 0.0 # far away is untouched + + def test_strokes_no_conjugate_when_disabled(self): + s = FourierStrokes(stamps=((KX, 0),), radius=3) + m = fourier_region_mask((N, N), [s], conjugate=False) + assert m[CY, CX + KX] == 1.0 + assert m[CY, CX - KX] == 0.0 + + def test_soft_edge_graded_rect_and_strokes(self): + rect = fourier_region_mask((N, N), [FourierRect(0, 0, 5, 5)], soft_px=3.0) + strokes = fourier_region_mask((N, N), [FourierStrokes(((0, 0),), 5)], soft_px=3.0) + for m in (rect, strokes): + assert ((m > 0.0) & (m < 1.0)).any() + + def test_region_from_dict_kinds(self): + assert isinstance(fourier_region_from_dict({"kind": "rect", "half_w": 2, "half_h": 3}), + FourierRect) + assert isinstance(fourier_region_from_dict({"kind": "paint", "stamps": [[1, 2]], "radius": 2}), + FourierStrokes) + # legacy dict without a kind reads as an ellipse + assert isinstance(fourier_region_from_dict({"dx": 1, "dy": 0, "rx": 2, "ry": 2}), + FourierEllipse) + + def test_mixed_regions_via_filter(self): + _b, _s, img = _scene() + regions = [ + FourierEllipse(dx=KX, dy=0, rx=3, ry=3), + FourierRect(dx=20, dy=10, half_w=2, half_h=2), + FourierStrokes(stamps=((-15, -5),), radius=2), + ] + out = inverse_fft_filter(img, regions, mode="remove_selected") + assert out.shape == img.shape and np.isfinite(out).all() + assert _bin_power(out) < 1e-6 * _bin_power(img) # the sine ellipse still removed + + # ─── inverse_fft_from_mask ─────────────────────────────────────────────────── class TestInverse: