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:
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user