import cv2 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/polygon mask drawing, undo/redo, and erase.""" _BRUSH_DEFAULT = 5 _ALPHA_DEFAULT = 40 _BRIGHTNESS_DEFAULT = 0 _CONTRAST_DEFAULT = 0 _GAMMA_DEFAULT = 100 _CLOSE_THRESHOLD = 15 # image-pixel distance to first vertex that closes a polygon 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._current_frame = frames[0] self.tool_mode = "brush" self._shapes: list[list[tuple]] = [] self._current_poly: list[tuple] = [] self._poly_artists: list = [] self._mouse_pos: tuple | None = None 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.btn_brush = QPushButton("Brush") self.btn_brush.setStyleSheet("background-color: #4488ff; color: white;") self.btn_polygon = QPushButton("Polygon") self.btn_fill = QPushButton("Fill") self.btn_fill.setEnabled(False) self.btn_del_shape = QPushButton("Del Shape") self.btn_del_shape.setEnabled(False) self.btn_cancel_poly = QPushButton("Cancel Current Poly") self.brush_slider = QSlider(Qt.Horizontal) self.brush_slider.setRange(2, 50) self.brush_slider.setValue(self._BRUSH_DEFAULT) self.brush_reset = QPushButton("↺") self.brush_reset.setFixedWidth(28) self.alpha_slider = QSlider(Qt.Horizontal) self.alpha_slider.setRange(0, 100) self.alpha_slider.setValue(self._ALPHA_DEFAULT) self.alpha_reset = QPushButton("↺") self.alpha_reset.setFixedWidth(28) self.brightness_slider = QSlider(Qt.Vertical) self.brightness_slider.setRange(-100, 100) self.brightness_slider.setValue(self._BRIGHTNESS_DEFAULT) self.brightness_reset = QPushButton("↺") self.brightness_reset.setFixedWidth(28) self.contrast_slider = QSlider(Qt.Vertical) self.contrast_slider.setRange(-100, 100) self.contrast_slider.setValue(self._CONTRAST_DEFAULT) self.contrast_reset = QPushButton("↺") self.contrast_reset.setFixedWidth(28) self.gamma_slider = QSlider(Qt.Vertical) self.gamma_slider.setRange(10, 300) self.gamma_slider.setValue(self._GAMMA_DEFAULT) self.gamma_reset = QPushButton("↺") self.gamma_reset.setFixedWidth(28) 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) self.btn_brush.clicked.connect(lambda: self.set_tool_mode("brush")) self.btn_polygon.clicked.connect(lambda: self.set_tool_mode("polygon")) self.btn_fill.clicked.connect(lambda: self.set_tool_mode("fill")) self.btn_del_shape.clicked.connect(self.delete_last_shape) self.btn_cancel_poly.clicked.connect(self.cancel_polygon) self.alpha_slider.valueChanged.connect(self.redraw) self.brightness_slider.valueChanged.connect(self._refresh_frame) self.contrast_slider.valueChanged.connect(self._refresh_frame) self.gamma_slider.valueChanged.connect(self._refresh_frame) self.brush_reset.clicked.connect( lambda: self.brush_slider.setValue(self._BRUSH_DEFAULT) ) self.alpha_reset.clicked.connect( lambda: self.alpha_slider.setValue(self._ALPHA_DEFAULT) ) self.brightness_reset.clicked.connect( lambda: self.brightness_slider.setValue(self._BRIGHTNESS_DEFAULT) ) self.contrast_reset.clicked.connect( lambda: self.contrast_slider.setValue(self._CONTRAST_DEFAULT) ) self.gamma_reset.clicked.connect( lambda: self.gamma_slider.setValue(self._GAMMA_DEFAULT) ) # ── 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._current_frame = frames[0] self._clear_poly_state() self.img_artist.set_data(self._apply_image_adjustments(frames[0])) self.set_title(title) self.redraw() def _clear_poly_state(self): self._shapes = [] self._current_poly = [] self._mouse_pos = None for a in self._poly_artists: a.remove() self._poly_artists = [] self._update_poly_buttons() # ── frame / title ────────────────────────────────────────────── def set_frame(self, frame): self._current_frame = frame self.img_artist.set_data(self._apply_image_adjustments(frame)) self.canvas.draw_idle() # ── image adjustments ────────────────────────────────────────── def _apply_image_adjustments(self, frame): img = frame.astype(np.float32) img += self.brightness_slider.value() c = self.contrast_slider.value() / 100.0 img = (1.0 + c) * (img - 128.0) + 128.0 np.clip(img, 0, 255, out=img) g = self.gamma_slider.value() / 100.0 img = (img / 255.0) ** (1.0 / g) * 255.0 return np.clip(img, 0, 255).astype(np.uint8) def _refresh_frame(self): if self._current_frame is not None: self.img_artist.set_data(self._apply_image_adjustments(self._current_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 set_mask(self, mask): """Replace the mask and push the previous state onto the undo stack.""" self.history.append(self.mask.copy()) self.redo_stack.clear() self.mask = mask self.redraw() def redraw(self): if self.mask_visible: alpha = self.alpha_slider.value() / 100.0 rgba = np.zeros((self.dh, self.dw, 4)) rgba[..., 1] = self.mask * 0.7 rgba[..., 3] = self.mask * alpha 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 if self.erase_mode: self.btn_erase.setText("Eraser ON") self.btn_erase.setStyleSheet("background-color: orange; color: black;") else: self.btn_erase.setText("Eraser") self.btn_erase.setStyleSheet("") def toggle_mask(self): self.mask_visible = not self.mask_visible if self.mask_visible: self.btn_mask.setText("Hide Mask") self.btn_mask.setStyleSheet("") else: self.btn_mask.setText("Show Mask") self.btn_mask.setStyleSheet("background-color: red; color: white;") 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() # ── tool mode ────────────────────────────────────────────────── def set_tool_mode(self, mode: str): self.tool_mode = mode active = "background-color: #4488ff; color: white;" self.btn_brush.setStyleSheet(active if mode == "brush" else "") self.btn_polygon.setStyleSheet(active if mode == "polygon" else "") self.btn_fill.setStyleSheet(active if mode == "fill" else "") if mode != "brush": self.brush_circle.set_visible(False) self.canvas.draw_idle() # ── polygon ops ──────────────────────────────────────────────── def _near_first(self, x: float, y: float) -> bool: if not self._current_poly: return False fx, fy = self._current_poly[0] return (x - fx) ** 2 + (y - fy) ** 2 <= self._CLOSE_THRESHOLD**2 def _update_poly_buttons(self): has = bool(self._shapes) self.btn_fill.setEnabled(has) self.btn_del_shape.setEnabled(has) def _draw_polygon_overlay(self, mouse_pos=None): for a in self._poly_artists: a.remove() self._poly_artists.clear() # Completed shapes — thick closed outline for shape in self._shapes: xs = [p[0] for p in shape] + [shape[0][0]] ys = [p[1] for p in shape] + [shape[0][1]] (line,) = self.ax.plot(xs, ys, color="cyan", linewidth=3, zorder=5) (dots,) = self.ax.plot( [p[0] for p in shape], [p[1] for p in shape], "o", color="cyan", markersize=4, zorder=6, ) self._poly_artists.extend([line, dots]) # In-progress polygon if self._current_poly: xs = [p[0] for p in self._current_poly] ys = [p[1] for p in self._current_poly] if len(self._current_poly) > 1: (edge,) = self.ax.plot(xs, ys, color="yellow", linewidth=1.5, zorder=5) self._poly_artists.append(edge) (verts,) = self.ax.plot(xs, ys, "o", color="yellow", markersize=5, zorder=6) # Red dot on first vertex as close-target indicator (first,) = self.ax.plot( [xs[0]], [ys[0]], "o", color="red", markersize=8, zorder=7 ) self._poly_artists.extend([verts, first]) # Rubber-band line from last vertex to cursor if mouse_pos: mx, my = mouse_pos near = len(self._current_poly) >= 3 and self._near_first(mx, my) clr = "lime" if near else "yellow" (rband,) = self.ax.plot( [xs[-1], mx], [ys[-1], my], "--", color=clr, linewidth=1, zorder=5 ) self._poly_artists.append(rband) self.canvas.draw_idle() def cancel_polygon(self): self._current_poly = [] self._draw_polygon_overlay(mouse_pos=self._mouse_pos) def delete_last_shape(self): if self._shapes: self._shapes.pop() self._update_poly_buttons() self._draw_polygon_overlay(mouse_pos=self._mouse_pos) def _fill_shape_at(self, x: float, y: float): if not self._shapes: return polys = [ np.array( [(int(round(px)), int(round(py))) for px, py in shape], dtype=np.int32 ) for shape in self._shapes ] # Find all shapes that contain the click point containing = [] for i, poly in enumerate(polys): poly_f32 = poly.reshape(-1, 1, 2).astype(np.float32) if cv2.pointPolygonTest(poly_f32, (x, y), False) >= 0: containing.append((i, poly)) if not containing: return # click was outside all shapes # Pick the innermost (smallest area) shape that contains the click containing.sort(key=lambda t: cv2.contourArea(t[1])) target_idx, target_poly = containing[0] self.history.append(self.mask.copy()) self.redo_stack.clear() temp = np.zeros((self.dh, self.dw), dtype=np.uint8) cv2.fillPoly(temp, [target_poly], 1) # Punch holes for any shapes completely inside the target target_f32 = target_poly.reshape(-1, 1, 2).astype(np.float32) for i, poly in enumerate(polys): if i == target_idx: continue cx = float(np.mean(poly[:, 0])) cy = float(np.mean(poly[:, 1])) if cv2.pointPolygonTest(target_f32, (cx, cy), False) > 0: cv2.fillPoly(temp, [poly], 0) self.mask |= temp 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) if self.tool_mode == "polygon": self._mouse_pos = None self._draw_polygon_overlay() else: self.canvas.draw_idle() # ── mouse events ─────────────────────────────────────────────── def _on_press(self, e): if e.xdata is None: return if self.tool_mode == "brush": self.drawing = True self.stamp(e.xdata, e.ydata) elif self.tool_mode == "polygon": self._handle_polygon_click(e) elif self.tool_mode == "fill" and e.button == 1: self._fill_shape_at(e.xdata, e.ydata) def _handle_polygon_click(self, e): if e.button == 3: # right-click: remove last vertex if self._current_poly: self._current_poly.pop() self._draw_polygon_overlay(mouse_pos=self._mouse_pos) return if e.button != 1: return x, y = e.xdata, e.ydata if len(self._current_poly) >= 3 and self._near_first(x, y): self._shapes.append(list(self._current_poly)) self._current_poly = [] self._update_poly_buttons() self._draw_polygon_overlay(mouse_pos=self._mouse_pos) else: self._current_poly.append((x, y)) self._draw_polygon_overlay(mouse_pos=self._mouse_pos) def _on_move(self, e): if self.tool_mode == "brush": self._update_brush_preview(e) if self.drawing: self.stamp(e.xdata, e.ydata) elif self.tool_mode == "polygon": self.brush_circle.set_visible(False) if e.inaxes == self.ax and e.xdata is not None: self._mouse_pos = (e.xdata, e.ydata) self._draw_polygon_overlay(mouse_pos=self._mouse_pos) else: self._mouse_pos = None self._draw_polygon_overlay() def _on_release(self, _): self.drawing = False