Add polygon drawing and click-to-fill tools

Introduce two new drawing modes alongside the existing brush:

- Polygon mode: left-click to place vertices connected by lines; right-click
  removes the last vertex; clicking near the first vertex (red dot) when >= 3
  points are placed closes the shape (bold cyan outline). Multiple shapes can
  coexist as canvas overlays. Cancel Current Poly discards the in-progress
  polygon; Del Shape removes the last completed shape.

- Fill mode: left-click inside any closed polygon to rasterise it onto the
  mask. Selects the innermost shape containing the click (smallest area via
  cv2.contourArea). Polygons whose centroid lies inside the target are punched
  out as holes. Each fill is a single undoable step in the mask history.

Also renames the Alpha slider label to Mask Alpha and removes the stale
Reload Saved reference from the README.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
2026-05-20 15:02:46 +02:00
parent d13ad1743a
commit 47432cec4f
3 changed files with 267 additions and 18 deletions

View File

@@ -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)

View File

@@ -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