diff --git a/README.md b/README.md index 4e0a9a5..6386027 100644 --- a/README.md +++ b/README.md @@ -49,7 +49,7 @@ python -m river_annotation_tool.annotation_script | `--data` | *(from config)* | Override `data_dir` from config | | `--out` | *(from config)* | Override `out_dir` from config | | `--clips` | *(from config)* | Override `clips_file` from config | -| `--clip` | *(first unannotated in list)* | Open a specific clip by stem name (e.g. `left_20230501`) | +| `--clip` | *(first unannotated in list)* | Open a specific clip by stem name | | `--extras` | off | Also save GIFs and extra PNGs (see Output section) | | `--no-skip` | off | Show already-annotated clips instead of skipping them | @@ -124,21 +124,58 @@ Copy `config/clips.example.txt` as a starting point. The window shows the video on the left (auto-playing) and the survey panel on the right. -### Mask drawing +### Tool modes + +Three drawing tools are available in the tool row. The active tool is highlighted in blue. + +| Tool | How to activate | Description | +|---|---|---| +| **Brush** | Click **Brush** | Click and drag to paint the mask with a circular brush (default) | +| **Polygon** | Click **Polygon** | Click to place vertices and build closed shapes; use **Fill** mode to commit them | +| **Fill** | Click **Fill** | Click inside a closed polygon to fill it onto the mask | + +### Brush tool | Action | How | |---|---| | Draw water mask | Click and drag on the video | | Erase mask | Toggle **Eraser** button (turns orange when active), then drag | | Brush preview | A white circle follows the cursor showing the current brush size | -| Adjust brush size | **Brush size** slider below the controls; click **↺** to reset | -| Undo last stroke | **Undo** | -| Undo 10 strokes | **Undo×10** | -| Redo | **Redo** — steps forward through undone strokes | +| Adjust brush size | **Brush size** slider; click **↺** to reset | + +### Polygon tool + +Polygons are drawn as overlays and do not affect the mask until you use **Fill** mode. + +| Action | How | +|---|---| +| Add vertex | Left-click on the canvas | +| Remove last vertex | Right-click | +| Close a shape | Left-click near the first vertex (red dot) when ≥ 3 vertices are placed; completed shapes turn bold cyan | +| Draw multiple shapes | Each closed shape is kept independently; draw as many as needed | +| Cancel in-progress polygon | **Cancel Current Poly** — discards the unfinished polygon, keeps completed shapes | +| Delete last completed shape | **Del Shape** | + +### Fill tool + +| Action | How | +|---|---| +| Fill a shape | Left-click anywhere inside a closed polygon; that shape's interior is painted onto the mask | +| Nested shapes | If a closed polygon lies entirely inside the target, its interior is left unfilled (acts as a hole) | +| Innermost shape | Clicking inside nested shapes always fills the innermost (smallest) polygon containing the click | +| Undo fill | **Undo** — each fill is a single undoable step | + +### Common mask actions + +| Action | How | +|---|---| +| Undo last action | **Undo** | +| Undo 10 actions | **Undo×10** | +| Redo | **Redo** | | Clear entire mask | **Clear** | | Toggle mask overlay | **Hide Mask / Show Mask** — button turns red when hidden; does not affect mask data | -| Mask transparency | **Alpha** slider below the controls; click **↺** to reset | -| Load mask from previous clip | **Load Prev Mask** — copies the saved mask of the previous clip in the list onto the current clip; the action is undoable with **Undo** | +| Mask transparency | **Mask Alpha** slider; click **↺** to reset | +| Load mask from previous clip | **Load Prev Mask** — copies the saved mask of the previous clip; undoable with **Undo** | ### Image display adjustments @@ -208,11 +245,19 @@ Up to `max_frames` frames are extracted from the video and scaled so the longest ### Mask drawing -The mask is a binary NumPy array matching the display frame size. Each brush stroke stamps a filled circle of the selected radius, setting pixels to 1 (draw) or 0 (erase). The history stack stores a copy of the mask before each stroke, enabling unlimited undo. On save the mask is resized to the original video resolution with nearest-neighbour interpolation and written as an 8-bit PNG (0 or 255). +The mask is a binary NumPy array matching the display frame size. + +**Brush:** each stroke stamps a filled circle of the selected radius, setting pixels to 1 (draw) or 0 (erase). + +**Polygon:** vertices are stored as a list of floating-point canvas coordinates. Multiple closed shapes can coexist. Completed shapes are rendered as cyan overlays on the canvas but do not touch the mask until a fill is applied. + +**Fill:** clicking inside a closed polygon rasterises it with `cv2.fillPoly` and ORs the result into the mask. Among all shapes containing the click, the innermost (smallest area, determined by `cv2.contourArea`) is selected as the fill target. Any polygon whose centroid lies inside the target is then punched out as a hole. + +Every mask-changing operation (brush stroke, fill) pushes the previous mask onto the undo stack before modifying it. On save the mask is resized to the original video resolution with nearest-neighbour interpolation and written as an 8-bit PNG (0 or 255). ### Resuming -When a clip is loaded that already has a saved `mask.png` and `metadata.json`, the mask is restored at display resolution and the survey answers are pre-filled. **Reload Saved** lets you revert to the last save at any point during the current session. +When a clip is loaded that already has a saved `mask.png` and `metadata.json`, the mask is restored at display resolution and the survey answers are pre-filled. ## Repository structure diff --git a/src/river_annotation_tool/annotator.py b/src/river_annotation_tool/annotator.py index 2ab0aa1..f8c3753 100644 --- a/src/river_annotation_tool/annotator.py +++ b/src/river_annotation_tool/annotator.py @@ -119,6 +119,16 @@ class Annotator(QMainWindow): for b in [self.btn_prev, btn_next, btn_skip, btn_load_prev_mask]: row1.addWidget(b) + row_tools = QHBoxLayout() + for b in [ + self.mc.btn_brush, + self.mc.btn_polygon, + self.mc.btn_fill, + self.mc.btn_del_shape, + self.mc.btn_cancel_poly, + ]: + row_tools.addWidget(b) + row2 = QHBoxLayout() for b in [ btn_clear, @@ -136,7 +146,7 @@ class Annotator(QMainWindow): row3.addWidget(self.mc.brush_reset) row4 = QHBoxLayout() - row4.addWidget(QLabel("Alpha")) + row4.addWidget(QLabel("Mask Alpha")) row4.addWidget(self.mc.alpha_slider) row4.addWidget(self.mc.alpha_reset) @@ -162,6 +172,7 @@ class Annotator(QMainWindow): left = QVBoxLayout() left.addLayout(canvas_row) left.addLayout(row1) + left.addLayout(row_tools) left.addLayout(row2) left.addLayout(row3) left.addLayout(row4) diff --git a/src/river_annotation_tool/mask_canvas.py b/src/river_annotation_tool/mask_canvas.py index c5d4e15..34020e5 100644 --- a/src/river_annotation_tool/mask_canvas.py +++ b/src/river_annotation_tool/mask_canvas.py @@ -1,3 +1,4 @@ +import cv2 import numpy as np from matplotlib.backends.backend_qtagg import FigureCanvasQTAgg as FigureCanvas from matplotlib.figure import Figure @@ -7,13 +8,14 @@ from PySide6.QtWidgets import QPushButton, QSlider class MaskCanvas: - """Matplotlib canvas with brush-based mask drawing, undo/redo, and erase.""" + """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 @@ -27,6 +29,12 @@ class MaskCanvas: 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() @@ -48,6 +56,15 @@ class MaskCanvas: 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) @@ -85,6 +102,11 @@ class MaskCanvas: 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) @@ -113,10 +135,20 @@ class MaskCanvas: 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 @@ -226,6 +258,131 @@ class MaskCanvas: 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: @@ -238,19 +395,55 @@ class MaskCanvas: def _on_axes_leave(self, _): self.brush_circle.set_visible(False) - self.canvas.draw_idle() + 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 - self.drawing = True - self.stamp(e.xdata, e.ydata) + 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): - self._update_brush_preview(e) - if self.drawing: - self.stamp(e.xdata, e.ydata) + 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