115 lines
4.2 KiB
Python
115 lines
4.2 KiB
Python
import numpy as np
|
|
from matplotlib.backends.backend_qtagg import FigureCanvasQTAgg as FigureCanvas
|
|
from matplotlib.figure import Figure
|
|
from PySide6.QtCore import Qt
|
|
from PySide6.QtWidgets import QPushButton, QSlider
|
|
|
|
|
|
class MaskCanvas:
|
|
"""Matplotlib canvas with brush-based mask drawing, undo, 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.erase_mode = False
|
|
self.drawing = False
|
|
|
|
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)
|
|
|
|
def _build_controls(self):
|
|
self.btn_erase = QPushButton("Eraser")
|
|
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.btn_erase.clicked.connect(self.toggle_erase)
|
|
|
|
# ── 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.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.redraw()
|
|
|
|
def redraw(self):
|
|
rgba = np.zeros((self.dh, self.dw, 4))
|
|
rgba[..., 1] = self.mask * 0.7
|
|
rgba[..., 3] = self.mask * 0.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.mask = self.history.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 stamp(self, x, y):
|
|
if x is None or y is None:
|
|
return
|
|
self.history.append(self.mask.copy())
|
|
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()
|
|
|
|
# ── 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):
|
|
if self.drawing:
|
|
self.stamp(e.xdata, e.ydata)
|
|
|
|
def _on_release(self, _):
|
|
self.drawing = False
|