clip-annotator/src/river_annotation_tool/mask_canvas.py

import numpy as np
from matplotlib.backends.backend_qtagg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.figure import Figure
from matplotlib.patches import Circle
from PySide6.QtCore import Qt
from PySide6.QtWidgets import QPushButton, QSlider


class MaskCanvas:
    """Matplotlib canvas with brush-based mask drawing, undo/redo, and erase."""

    def __init__(self, frames, dh: int, dw: int):
        self.dh = dh
        self.dw = dw

        self.mask = np.zeros((dh, dw), dtype=np.uint8)
        self.history: list[np.ndarray] = []
        self.redo_stack: list[np.ndarray] = []
        self.erase_mode = False
        self.drawing = False
        self.mask_visible = True

        self._build_figure(frames)
        self._build_controls()
        self._connect_events()

    def _build_figure(self, frames):
        self.fig = Figure()
        self.canvas = FigureCanvas(self.fig)
        self.ax = self.fig.add_subplot(111)
        self.ax.axis("off")
        self.img_artist = self.ax.imshow(frames[0])
        self.mask_artist = self.ax.imshow(np.zeros((self.dh, self.dw, 4)))
        self.title_text = self.ax.set_title("", fontsize=10, pad=4)
        self.brush_circle = Circle(
            (0, 0), radius=5, fill=False, color="white", linewidth=1.5, visible=False
        )
        self.ax.add_patch(self.brush_circle)

    def _build_controls(self):
        self.btn_erase = QPushButton("Eraser")
        self.btn_mask = QPushButton("Hide Mask")
        self.brush_slider = QSlider(Qt.Horizontal)
        self.brush_slider.setRange(2, 50)
        self.brush_slider.setValue(5)

    def _connect_events(self):
        self.canvas.mpl_connect("button_press_event", self._on_press)
        self.canvas.mpl_connect("motion_notify_event", self._on_move)
        self.canvas.mpl_connect("button_release_event", self._on_release)
        self.canvas.mpl_connect("axes_leave_event", self._on_axes_leave)
        self.btn_erase.clicked.connect(self.toggle_erase)
        self.btn_mask.clicked.connect(self.toggle_mask)

    # ── clip transition ────────────────────────────────────────────
    def load_clip(self, frames, dh: int, dw: int, mask=None, title: str = ""):
        self.dh = dh
        self.dw = dw
        self.mask = mask if mask is not None else np.zeros((dh, dw), dtype=np.uint8)
        self.history = []
        self.redo_stack = []
        self.img_artist.set_data(frames[0])
        self.set_title(title)
        self.redraw()

    # ── frame / title ──────────────────────────────────────────────
    def set_frame(self, frame):
        self.img_artist.set_data(frame)
        self.canvas.draw_idle()

    def set_title(self, text: str):
        self.title_text.set_text(text)

    # ── mask ops ───────────────────────────────────────────────────
    def reset(self, mask=None):
        self.mask = (
            mask if mask is not None else np.zeros((self.dh, self.dw), dtype=np.uint8)
        )
        self.history = []
        self.redo_stack = []
        self.redraw()

    def redraw(self):
        if self.mask_visible:
            rgba = np.zeros((self.dh, self.dw, 4))
            rgba[..., 1] = self.mask * 0.7
            rgba[..., 3] = self.mask * 0.4
        else:
            rgba = np.zeros((self.dh, self.dw, 4))
        self.mask_artist.set_data(rgba)
        self.canvas.draw_idle()

    def clear(self):
        self.mask[:] = 0
        self.redraw()

    def undo(self):
        if self.history:
            self.redo_stack.append(self.mask.copy())
            self.mask = self.history.pop()
            self.redraw()

    def undo10(self):
        for _ in range(10):
            if not self.history:
                break
            self.redo_stack.append(self.mask.copy())
            self.mask = self.history.pop()
        self.redraw()

    def redo(self):
        if self.redo_stack:
            self.history.append(self.mask.copy())
            self.mask = self.redo_stack.pop()
            self.redraw()

    def toggle_erase(self):
        self.erase_mode = not self.erase_mode
        self.btn_erase.setText("Eraser ON" if self.erase_mode else "Eraser")

    def toggle_mask(self):
        self.mask_visible = not self.mask_visible
        self.btn_mask.setText("Show Mask" if not self.mask_visible else "Hide Mask")
        self.redraw()

    def stamp(self, x, y):
        if x is None or y is None:
            return
        self.history.append(self.mask.copy())
        self.redo_stack.clear()
        r = self.brush_slider.value()
        ix, iy = int(x), int(y)
        y0, y1 = max(0, iy - r), min(self.dh, iy + r + 1)
        x0, x1 = max(0, ix - r), min(self.dw, ix + r + 1)
        Y, X = np.ogrid[y0:y1, x0:x1]
        circle = (X - ix) ** 2 + (Y - iy) ** 2 <= r**2
        self.mask[y0:y1, x0:x1][circle] = 0 if self.erase_mode else 1
        self.redraw()

    # ── brush preview ──────────────────────────────────────────────
    def _update_brush_preview(self, e):
        if e.inaxes == self.ax and e.xdata is not None:
            self.brush_circle.center = (e.xdata, e.ydata)
            self.brush_circle.set_radius(self.brush_slider.value())
            self.brush_circle.set_visible(True)
        else:
            self.brush_circle.set_visible(False)
        self.canvas.draw_idle()

    def _on_axes_leave(self, _):
        self.brush_circle.set_visible(False)
        self.canvas.draw_idle()

    # ── mouse events ───────────────────────────────────────────────
    def _on_press(self, e):
        if e.xdata is None:
            return
        self.drawing = True
        self.stamp(e.xdata, e.ydata)

    def _on_move(self, e):
        self._update_brush_preview(e)
        if self.drawing:
            self.stamp(e.xdata, e.ydata)

    def _on_release(self, _):
        self.drawing = False