Made project river-agnostic

src/clip_annotator/mask_canvas.py

@@ -0,0 +1,452 @@
+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 = 15
+    _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(figsize=(self.dw / 80, self.dh / 80))
+        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)
+        self.ax.autoscale(False)  # prevent polygon plot() calls from expanding the view
+
+    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.ax.set_xlim(-0.5, dw - 0.5)
+        self.ax.set_ylim(dh - 0.5, -0.5)
+        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