1201 lines
44 KiB
Python
1201 lines
44 KiB
Python
from __future__ import annotations
|
|
|
|
import collections
|
|
import enum
|
|
from typing import Literal
|
|
|
|
import imgviz
|
|
import numpy as np
|
|
import osam
|
|
from loguru import logger
|
|
from PyQt5 import QtCore
|
|
from PyQt5 import QtGui
|
|
from PyQt5 import QtWidgets
|
|
from PyQt5.QtCore import QPoint
|
|
from PyQt5.QtCore import QPointF
|
|
from PyQt5.QtCore import Qt
|
|
|
|
import labelme.utils
|
|
from labelme._automation import polygon_from_mask
|
|
from labelme.shape import Shape
|
|
|
|
from .download import download_ai_model
|
|
|
|
# TODO(unknown):
|
|
# - [maybe] Find optimal epsilon value.
|
|
|
|
|
|
CURSOR_DEFAULT = Qt.ArrowCursor
|
|
CURSOR_POINT = Qt.PointingHandCursor
|
|
CURSOR_DRAW = Qt.CrossCursor
|
|
CURSOR_MOVE = Qt.ClosedHandCursor
|
|
CURSOR_GRAB = Qt.OpenHandCursor
|
|
|
|
MOVE_SPEED = 5.0
|
|
|
|
|
|
class CanvasMode(enum.Enum):
|
|
CREATE = enum.auto()
|
|
EDIT = enum.auto()
|
|
|
|
|
|
class Canvas(QtWidgets.QWidget):
|
|
pixmap: QtGui.QPixmap
|
|
_cursor: QtCore.Qt.CursorShape
|
|
shapes: list[Shape]
|
|
shapesBackups: list[list[Shape]]
|
|
movingShape: bool
|
|
selectedShapes: list[Shape]
|
|
selectedShapesCopy: list[Shape]
|
|
current: Shape | None
|
|
hShape: Shape | None
|
|
prevhShape: Shape | None
|
|
hVertex: int | None
|
|
prevhVertex: int | None
|
|
hEdge: int | None
|
|
prevhEdge: int | None
|
|
|
|
zoomRequest = QtCore.pyqtSignal(int, QPointF)
|
|
scrollRequest = QtCore.pyqtSignal(int, int)
|
|
newShape = QtCore.pyqtSignal()
|
|
selectionChanged = QtCore.pyqtSignal(list)
|
|
shapeMoved = QtCore.pyqtSignal()
|
|
drawingPolygon = QtCore.pyqtSignal(bool)
|
|
vertexSelected = QtCore.pyqtSignal(bool)
|
|
mouseMoved = QtCore.pyqtSignal(QPointF)
|
|
statusUpdated = QtCore.pyqtSignal(str)
|
|
|
|
mode: CanvasMode = CanvasMode.EDIT
|
|
|
|
# polygon, rectangle, line, or point
|
|
_createMode = "polygon"
|
|
|
|
_fill_drawing = False
|
|
|
|
prevPoint: QPointF
|
|
prevMovePoint: QPointF
|
|
offsets: tuple[QPointF, QPointF]
|
|
|
|
_dragging_start_pos: QPointF
|
|
_is_dragging: bool
|
|
_is_dragging_enabled: bool
|
|
|
|
_ai_model_name: str = "sam2:latest"
|
|
_ai_model_cache: osam.types.Model | None = None
|
|
_ai_image_embedding_cache: collections.deque[tuple[str, osam.types.ImageEmbedding]]
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
self.epsilon = kwargs.pop("epsilon", 10.0)
|
|
self.double_click = kwargs.pop("double_click", "close")
|
|
if self.double_click not in [None, "close"]:
|
|
raise ValueError(
|
|
f"Unexpected value for double_click event: {self.double_click}"
|
|
)
|
|
self.num_backups = kwargs.pop("num_backups", 10)
|
|
self._crosshair = kwargs.pop(
|
|
"crosshair",
|
|
{
|
|
"polygon": False,
|
|
"rectangle": True,
|
|
"circle": False,
|
|
"line": False,
|
|
"point": False,
|
|
"linestrip": False,
|
|
"ai_polygon": False,
|
|
"ai_mask": False,
|
|
},
|
|
)
|
|
super().__init__(*args, **kwargs)
|
|
|
|
self.resetState()
|
|
|
|
# self.line represents:
|
|
# - createMode == 'polygon': edge from last point to current
|
|
# - createMode == 'rectangle': diagonal line of the rectangle
|
|
# - createMode == 'line': the line
|
|
# - createMode == 'point': the point
|
|
self.line = Shape()
|
|
self.prevPoint = QPointF()
|
|
self.prevMovePoint = QPointF()
|
|
self.offsets = QPointF(), QPointF()
|
|
self.scale = 1.0
|
|
self._ai_image_embedding_cache = collections.deque(maxlen=3)
|
|
self.visible = {}
|
|
self._hideBackround = False
|
|
self.hideBackround = False
|
|
self.snapping = True
|
|
self.hShapeIsSelected = False
|
|
self._painter = QtGui.QPainter()
|
|
self._dragging_start_pos = QPointF()
|
|
self._is_dragging = False
|
|
self._is_dragging_enabled = False
|
|
# Menus:
|
|
# 0: right-click without selection and dragging of shapes
|
|
# 1: right-click with selection and dragging of shapes
|
|
self.menus = (QtWidgets.QMenu(), QtWidgets.QMenu())
|
|
# Set widget options.
|
|
self.setMouseTracking(True)
|
|
self.setFocusPolicy(Qt.WheelFocus)
|
|
|
|
def fillDrawing(self):
|
|
return self._fill_drawing
|
|
|
|
def setFillDrawing(self, value):
|
|
self._fill_drawing = value
|
|
|
|
@property
|
|
def createMode(self):
|
|
return self._createMode
|
|
|
|
@createMode.setter
|
|
def createMode(self, value):
|
|
if value not in [
|
|
"polygon",
|
|
"rectangle",
|
|
"circle",
|
|
"line",
|
|
"point",
|
|
"linestrip",
|
|
"ai_polygon",
|
|
"ai_mask",
|
|
]:
|
|
raise ValueError(f"Unsupported createMode: {value}")
|
|
self._createMode = value
|
|
|
|
def set_ai_model_name(self, model_name: str) -> None:
|
|
logger.debug("Setting AI model to {!r}", model_name)
|
|
self._ai_model_name = model_name
|
|
|
|
def _get_ai_model(self) -> osam.types.Model:
|
|
if self._ai_model_cache and self._ai_model_cache.name == self._ai_model_name:
|
|
return self._ai_model_cache
|
|
|
|
model_type = osam.apis.get_model_type_by_name(self._ai_model_name)
|
|
|
|
self._ai_model_cache = model_type()
|
|
return self._ai_model_cache
|
|
|
|
def _get_ai_image_embedding(self) -> osam.types.ImageEmbedding:
|
|
qimage: QtGui.QImage = self.pixmap.toImage()
|
|
|
|
def pixmap_hash() -> int:
|
|
bits = qimage.constBits()
|
|
if bits is None:
|
|
return hash(None)
|
|
return hash(bits.asstring(qimage.sizeInBytes()))
|
|
|
|
cache_key: str = f"{self._ai_model_name}_{pixmap_hash()}"
|
|
key: str
|
|
image_embedding: osam.types.ImageEmbedding
|
|
for key, image_embedding in self._ai_image_embedding_cache:
|
|
if key == cache_key:
|
|
return image_embedding
|
|
|
|
image: np.ndarray = labelme.utils.img_qt_to_arr(img_qt=qimage)
|
|
image_embedding = self._get_ai_model().encode_image(image=imgviz.asrgb(image))
|
|
self._ai_image_embedding_cache.append((cache_key, image_embedding))
|
|
logger.debug("cached image embedding for key: {!r}", cache_key)
|
|
return image_embedding
|
|
|
|
def storeShapes(self):
|
|
shapesBackup = []
|
|
for shape in self.shapes:
|
|
shapesBackup.append(shape.copy())
|
|
if len(self.shapesBackups) > self.num_backups:
|
|
self.shapesBackups = self.shapesBackups[-self.num_backups - 1 :]
|
|
self.shapesBackups.append(shapesBackup)
|
|
|
|
@property
|
|
def isShapeRestorable(self):
|
|
# We save the state AFTER each edit (not before) so for an
|
|
# edit to be undoable, we expect the CURRENT and the PREVIOUS state
|
|
# to be in the undo stack.
|
|
if len(self.shapesBackups) < 2:
|
|
return False
|
|
return True
|
|
|
|
def restoreShape(self):
|
|
# This does _part_ of the job of restoring shapes.
|
|
# The complete process is also done in app.py::undoShapeEdit
|
|
# and app.py::loadShapes and our own Canvas::loadShapes function.
|
|
if not self.isShapeRestorable:
|
|
return
|
|
self.shapesBackups.pop() # latest
|
|
|
|
# The application will eventually call Canvas.loadShapes which will
|
|
# push this right back onto the stack.
|
|
shapesBackup = self.shapesBackups.pop()
|
|
self.shapes = shapesBackup
|
|
self.selectedShapes = []
|
|
for shape in self.shapes:
|
|
shape.selected = False
|
|
self.update()
|
|
|
|
def enterEvent(self, a0: QtCore.QEvent) -> None:
|
|
self.overrideCursor(self._cursor)
|
|
self._update_status()
|
|
|
|
def leaveEvent(self, a0: QtCore.QEvent) -> None:
|
|
self.unHighlight()
|
|
self.restoreCursor()
|
|
self._update_status()
|
|
|
|
def focusOutEvent(self, a0: QtGui.QFocusEvent) -> None:
|
|
self.restoreCursor()
|
|
self._update_status()
|
|
|
|
def isVisible(self, shape): # type: ignore[override]
|
|
return self.visible.get(shape, True)
|
|
|
|
def drawing(self):
|
|
return self.mode == CanvasMode.CREATE
|
|
|
|
def editing(self):
|
|
return self.mode == CanvasMode.EDIT
|
|
|
|
def setEditing(self, value=True):
|
|
self.mode = CanvasMode.EDIT if value else CanvasMode.CREATE
|
|
if self.mode == CanvasMode.EDIT:
|
|
# CREATE -> EDIT
|
|
self.repaint() # clear crosshair
|
|
else:
|
|
# EDIT -> CREATE
|
|
self.unHighlight()
|
|
self.deSelectShape()
|
|
|
|
def unHighlight(self):
|
|
if self.hShape:
|
|
self.hShape.highlightClear()
|
|
self.update()
|
|
self.prevhShape = self.hShape
|
|
self.prevhVertex = self.hVertex
|
|
self.prevhEdge = self.hEdge
|
|
self.hShape = self.hVertex = self.hEdge = None
|
|
|
|
def selectedVertex(self):
|
|
return self.hVertex is not None
|
|
|
|
def selectedEdge(self):
|
|
return self.hEdge is not None
|
|
|
|
def _update_status(self, extra_messages: list[str] | None = None) -> None:
|
|
messages: list[str] = []
|
|
if self.drawing():
|
|
messages.append(self.tr("Creating %r") % self.createMode)
|
|
messages.append(self._get_create_mode_message())
|
|
if self.current:
|
|
messages.append(self.tr("ESC to cancel"))
|
|
if self.canCloseShape():
|
|
messages.append(self.tr("Enter or Space to finalize"))
|
|
else:
|
|
assert self.editing()
|
|
messages.append(self.tr("Editing shapes"))
|
|
if extra_messages:
|
|
messages.extend(extra_messages)
|
|
self.statusUpdated.emit(" • ".join(messages))
|
|
|
|
def _get_create_mode_message(self) -> str:
|
|
assert self.drawing()
|
|
isNew: bool = self.current is None
|
|
if self.createMode == "ai_polygon":
|
|
return self.tr(
|
|
"Click points to include or Shift+Click to exclude for ai_polygon"
|
|
)
|
|
if self.createMode == "ai_mask":
|
|
return self.tr(
|
|
"Click points to include or Shift+Click to exclude for ai_mask"
|
|
)
|
|
if self.createMode == "line":
|
|
if isNew:
|
|
return self.tr("Click start point for line")
|
|
else:
|
|
return self.tr("Click end point for line")
|
|
if self.createMode == "linestrip":
|
|
if isNew:
|
|
return self.tr("Click start point for linestrip")
|
|
else:
|
|
return self.tr(
|
|
"Click next point or finish by Ctrl/Cmd+Click for linestrip"
|
|
)
|
|
if self.createMode == "circle":
|
|
if isNew:
|
|
return self.tr("Click center point for circle")
|
|
else:
|
|
return self.tr("Click point on circumference for circle")
|
|
if self.createMode == "rectangle":
|
|
if isNew:
|
|
return self.tr("Click first corner for rectangle")
|
|
else:
|
|
return self.tr("Click opposite corner for rectangle")
|
|
return self.tr("Click to add point")
|
|
|
|
def mouseMoveEvent(self, a0: QtGui.QMouseEvent) -> None:
|
|
"""Update line with last point and current coordinates."""
|
|
try:
|
|
pos = self.transformPos(a0.localPos())
|
|
except AttributeError:
|
|
return
|
|
|
|
self.mouseMoved.emit(pos)
|
|
|
|
self.prevMovePoint = pos
|
|
|
|
is_shift_pressed = a0.modifiers() & Qt.ShiftModifier
|
|
|
|
if self._is_dragging:
|
|
self.overrideCursor(CURSOR_GRAB)
|
|
delta: QPointF = pos - self._dragging_start_pos
|
|
self.scrollRequest.emit(int(delta.x()), Qt.Horizontal)
|
|
self.scrollRequest.emit(int(delta.y()), Qt.Vertical)
|
|
return
|
|
|
|
# Polygon drawing.
|
|
if self.drawing():
|
|
if self.createMode in ["ai_polygon", "ai_mask"]:
|
|
self.line.shape_type = "points"
|
|
else:
|
|
self.line.shape_type = self.createMode
|
|
|
|
self.overrideCursor(CURSOR_DRAW)
|
|
if not self.current:
|
|
self.repaint() # draw crosshair
|
|
self._update_status()
|
|
return
|
|
|
|
if self.outOfPixmap(pos):
|
|
# Don't allow the user to draw outside the pixmap.
|
|
# Project the point to the pixmap's edges.
|
|
pos = self.intersectionPoint(self.current[-1], pos)
|
|
elif (
|
|
self.snapping
|
|
and len(self.current) > 1
|
|
and self.createMode == "polygon"
|
|
and self.closeEnough(pos, self.current[0])
|
|
):
|
|
# Attract line to starting point and
|
|
# colorise to alert the user.
|
|
pos = self.current[0]
|
|
self.overrideCursor(CURSOR_POINT)
|
|
self.current.highlightVertex(0, Shape.NEAR_VERTEX)
|
|
if self.createMode in ["polygon", "linestrip"]:
|
|
self.line.points = [self.current[-1], pos]
|
|
self.line.point_labels = [1, 1]
|
|
elif self.createMode in ["ai_polygon", "ai_mask"]:
|
|
self.line.points = [self.current.points[-1], pos]
|
|
self.line.point_labels = [
|
|
self.current.point_labels[-1],
|
|
0 if is_shift_pressed else 1,
|
|
]
|
|
elif self.createMode == "rectangle":
|
|
self.line.points = [self.current[0], pos]
|
|
self.line.point_labels = [1, 1]
|
|
self.line.close()
|
|
elif self.createMode == "circle":
|
|
self.line.points = [self.current[0], pos]
|
|
self.line.point_labels = [1, 1]
|
|
self.line.shape_type = "circle"
|
|
elif self.createMode == "line":
|
|
self.line.points = [self.current[0], pos]
|
|
self.line.point_labels = [1, 1]
|
|
self.line.close()
|
|
elif self.createMode == "point":
|
|
self.line.points = [self.current[0]]
|
|
self.line.point_labels = [1]
|
|
self.line.close()
|
|
assert len(self.line.points) == len(self.line.point_labels)
|
|
self.repaint()
|
|
self.current.highlightClear()
|
|
self._update_status()
|
|
return
|
|
|
|
# Polygon copy moving.
|
|
if Qt.RightButton & a0.buttons():
|
|
if self.selectedShapesCopy and self.prevPoint is not None:
|
|
self.overrideCursor(CURSOR_MOVE)
|
|
self.boundedMoveShapes(self.selectedShapesCopy, pos)
|
|
self.repaint()
|
|
elif self.selectedShapes:
|
|
self.selectedShapesCopy = [s.copy() for s in self.selectedShapes]
|
|
self.repaint()
|
|
self._update_status()
|
|
return
|
|
|
|
# Polygon/Vertex moving.
|
|
if Qt.LeftButton & a0.buttons():
|
|
if self.selectedVertex():
|
|
self.boundedMoveVertex(pos)
|
|
self.repaint()
|
|
self.movingShape = True
|
|
elif self.selectedShapes and self.prevPoint is not None:
|
|
self.overrideCursor(CURSOR_MOVE)
|
|
self.boundedMoveShapes(self.selectedShapes, pos)
|
|
self.repaint()
|
|
self.movingShape = True
|
|
return
|
|
|
|
# Just hovering over the canvas, 2 possibilities:
|
|
# - Highlight shapes
|
|
# - Highlight vertex
|
|
# Update shape/vertex fill and tooltip value accordingly.
|
|
status_messages: list[str] = []
|
|
for shape in ([self.hShape] if self.hShape else []) + [
|
|
s for s in reversed(self.shapes) if self.isVisible(s) and s != self.hShape
|
|
]:
|
|
# Look for a nearby vertex to highlight. If that fails,
|
|
# check if we happen to be inside a shape.
|
|
index = shape.nearestVertex(pos, self.epsilon)
|
|
index_edge = shape.nearestEdge(pos, self.epsilon)
|
|
if index is not None:
|
|
if self.selectedVertex() and self.hShape:
|
|
self.hShape.highlightClear()
|
|
self.prevhVertex = self.hVertex = index
|
|
self.prevhShape = self.hShape = shape
|
|
self.prevhEdge = self.hEdge
|
|
self.hEdge = None
|
|
shape.highlightVertex(index, shape.MOVE_VERTEX)
|
|
self.overrideCursor(CURSOR_POINT)
|
|
status_messages.append(self.tr("Click & drag to move point"))
|
|
if shape.canRemovePoint():
|
|
status_messages.append(
|
|
self.tr("ALT + SHIFT + Click to delete point")
|
|
)
|
|
self.update()
|
|
break
|
|
elif index_edge is not None and shape.canAddPoint():
|
|
if self.selectedVertex() and self.hShape:
|
|
self.hShape.highlightClear()
|
|
self.prevhVertex = self.hVertex
|
|
self.hVertex = None
|
|
self.prevhShape = self.hShape = shape
|
|
self.prevhEdge = self.hEdge = index_edge
|
|
self.overrideCursor(CURSOR_POINT)
|
|
status_messages.append(self.tr("ALT + Click to create point on shape"))
|
|
self.update()
|
|
break
|
|
elif shape.containsPoint(pos):
|
|
if self.selectedVertex() and self.hShape:
|
|
self.hShape.highlightClear()
|
|
self.prevhVertex = self.hVertex
|
|
self.hVertex = None
|
|
self.prevhShape = self.hShape = shape
|
|
self.prevhEdge = self.hEdge
|
|
self.hEdge = None
|
|
status_messages.extend(
|
|
[
|
|
self.tr("Click & drag to move shape"),
|
|
self.tr("Right-click & drag to copy shape"),
|
|
]
|
|
)
|
|
self.overrideCursor(CURSOR_GRAB)
|
|
self.update()
|
|
break
|
|
else: # Nothing found, clear highlights, reset state.
|
|
self.restoreCursor()
|
|
self.unHighlight()
|
|
self.vertexSelected.emit(self.hVertex is not None)
|
|
self._update_status(extra_messages=status_messages)
|
|
|
|
def addPointToEdge(self):
|
|
shape = self.prevhShape
|
|
index = self.prevhEdge
|
|
point = self.prevMovePoint
|
|
if shape is None or index is None or point is None:
|
|
return
|
|
shape.insertPoint(index, point)
|
|
shape.highlightVertex(index, shape.MOVE_VERTEX)
|
|
self.hShape = shape
|
|
self.hVertex = index
|
|
self.hEdge = None
|
|
self.movingShape = True
|
|
|
|
def removeSelectedPoint(self):
|
|
shape = self.prevhShape
|
|
index = self.prevhVertex
|
|
if shape is None or index is None:
|
|
return
|
|
shape.removePoint(index)
|
|
shape.highlightClear()
|
|
self.hShape = shape
|
|
self.prevhVertex = None
|
|
self.movingShape = True # Save changes
|
|
|
|
def mousePressEvent(self, a0: QtGui.QMouseEvent) -> None:
|
|
pos: QPointF = self.transformPos(a0.localPos())
|
|
|
|
is_shift_pressed = a0.modifiers() & Qt.ShiftModifier
|
|
|
|
if a0.button() == Qt.LeftButton:
|
|
if self.drawing():
|
|
if self.current:
|
|
# Add point to existing shape.
|
|
if self.createMode == "polygon":
|
|
self.current.addPoint(self.line[1])
|
|
self.line[0] = self.current[-1]
|
|
if self.current.isClosed():
|
|
self.finalise()
|
|
elif self.createMode in ["rectangle", "circle", "line"]:
|
|
assert len(self.current.points) == 1
|
|
self.current.points = self.line.points
|
|
self.finalise()
|
|
elif self.createMode == "linestrip":
|
|
self.current.addPoint(self.line[1])
|
|
self.line[0] = self.current[-1]
|
|
if int(a0.modifiers()) == Qt.ControlModifier:
|
|
self.finalise()
|
|
elif self.createMode in ["ai_polygon", "ai_mask"]:
|
|
self.current.addPoint(
|
|
self.line.points[1],
|
|
label=self.line.point_labels[1],
|
|
)
|
|
self.line.points[0] = self.current.points[-1]
|
|
self.line.point_labels[0] = self.current.point_labels[-1]
|
|
if a0.modifiers() & Qt.ControlModifier:
|
|
self.finalise()
|
|
elif not self.outOfPixmap(pos):
|
|
if self.createMode in ["ai_polygon", "ai_mask"]:
|
|
if not download_ai_model(
|
|
model_name=self._ai_model_name, parent=self
|
|
):
|
|
return
|
|
|
|
# Create new shape.
|
|
self.current = Shape(
|
|
shape_type="points"
|
|
if self.createMode in ["ai_polygon", "ai_mask"]
|
|
else self.createMode
|
|
)
|
|
self.current.addPoint(pos, label=0 if is_shift_pressed else 1)
|
|
if self.createMode == "point":
|
|
self.finalise()
|
|
elif (
|
|
self.createMode in ["ai_polygon", "ai_mask"]
|
|
and a0.modifiers() & Qt.ControlModifier
|
|
):
|
|
self.finalise()
|
|
else:
|
|
if self.createMode == "circle":
|
|
self.current.shape_type = "circle"
|
|
self.line.points = [pos, pos]
|
|
if (
|
|
self.createMode in ["ai_polygon", "ai_mask"]
|
|
and is_shift_pressed
|
|
):
|
|
self.line.point_labels = [0, 0]
|
|
else:
|
|
self.line.point_labels = [1, 1]
|
|
self.setHiding()
|
|
self.drawingPolygon.emit(True)
|
|
self.update()
|
|
elif self.editing():
|
|
if self.selectedEdge() and a0.modifiers() == Qt.AltModifier:
|
|
self.addPointToEdge()
|
|
elif self.selectedVertex() and a0.modifiers() == (
|
|
Qt.AltModifier | Qt.ShiftModifier
|
|
):
|
|
self.removeSelectedPoint()
|
|
|
|
group_mode = int(a0.modifiers()) == Qt.ControlModifier
|
|
self.selectShapePoint(pos, multiple_selection_mode=group_mode)
|
|
self.prevPoint = pos
|
|
self.repaint()
|
|
elif a0.button() == Qt.RightButton and self.editing():
|
|
group_mode = int(a0.modifiers()) == Qt.ControlModifier
|
|
if not self.selectedShapes or (
|
|
self.hShape is not None and self.hShape not in self.selectedShapes
|
|
):
|
|
self.selectShapePoint(pos, multiple_selection_mode=group_mode)
|
|
self.repaint()
|
|
self.prevPoint = pos
|
|
elif a0.button() == Qt.MiddleButton and self._is_dragging_enabled:
|
|
self.overrideCursor(CURSOR_GRAB)
|
|
self._dragging_start_pos = pos
|
|
self._is_dragging = True
|
|
self._update_status()
|
|
|
|
def mouseReleaseEvent(self, a0: QtGui.QMouseEvent) -> None:
|
|
if a0.button() == Qt.RightButton:
|
|
menu = self.menus[len(self.selectedShapesCopy) > 0]
|
|
self.restoreCursor()
|
|
if not menu.exec_(self.mapToGlobal(a0.pos())) and self.selectedShapesCopy: # type: ignore
|
|
# Cancel the move by deleting the shadow copy.
|
|
self.selectedShapesCopy = []
|
|
self.repaint()
|
|
elif a0.button() == Qt.LeftButton:
|
|
if self.editing():
|
|
if (
|
|
self.hShape is not None
|
|
and self.hShapeIsSelected
|
|
and not self.movingShape
|
|
):
|
|
self.selectionChanged.emit(
|
|
[x for x in self.selectedShapes if x != self.hShape]
|
|
)
|
|
elif a0.button() == Qt.MiddleButton:
|
|
self._is_dragging = False
|
|
self.restoreCursor()
|
|
|
|
if self.movingShape and self.hShape:
|
|
index = self.shapes.index(self.hShape)
|
|
if self.shapesBackups[-1][index].points != self.shapes[index].points:
|
|
self.storeShapes()
|
|
self.shapeMoved.emit()
|
|
|
|
self.movingShape = False
|
|
self._update_status()
|
|
|
|
def endMove(self, copy):
|
|
assert self.selectedShapes and self.selectedShapesCopy
|
|
assert len(self.selectedShapesCopy) == len(self.selectedShapes)
|
|
if copy:
|
|
for i, shape in enumerate(self.selectedShapesCopy):
|
|
self.shapes.append(shape)
|
|
self.selectedShapes[i].selected = False
|
|
self.selectedShapes[i] = shape
|
|
else:
|
|
for i, shape in enumerate(self.selectedShapesCopy):
|
|
self.selectedShapes[i].points = shape.points
|
|
self.selectedShapesCopy = []
|
|
self.repaint()
|
|
self.storeShapes()
|
|
return True
|
|
|
|
def hideBackroundShapes(self, value):
|
|
self.hideBackround = value
|
|
if self.selectedShapes:
|
|
# Only hide other shapes if there is a current selection.
|
|
# Otherwise the user will not be able to select a shape.
|
|
self.setHiding(True)
|
|
self.update()
|
|
|
|
def setHiding(self, enable=True):
|
|
self._hideBackround = self.hideBackround if enable else False
|
|
|
|
def canCloseShape(self) -> bool:
|
|
if not self.drawing():
|
|
return False
|
|
if not self.current:
|
|
return False
|
|
if self.createMode in ["ai_polygon", "ai_mask"]:
|
|
return True
|
|
if self.createMode == "linestrip":
|
|
return len(self.current) >= 2
|
|
return len(self.current) >= 3
|
|
|
|
def mouseDoubleClickEvent(self, a0: QtGui.QMouseEvent) -> None:
|
|
if self.double_click != "close":
|
|
return
|
|
|
|
if self.canCloseShape():
|
|
self.finalise()
|
|
|
|
def selectShapes(self, shapes):
|
|
self.setHiding()
|
|
self.selectionChanged.emit(shapes)
|
|
self.update()
|
|
|
|
def selectShapePoint(self, point, multiple_selection_mode):
|
|
"""Select the first shape created which contains this point."""
|
|
if self.hVertex is not None:
|
|
assert self.hShape is not None
|
|
self.hShape.highlightVertex(i=self.hVertex, action=self.hShape.MOVE_VERTEX)
|
|
else:
|
|
shape: Shape
|
|
for shape in reversed(self.shapes):
|
|
if self.isVisible(shape) and shape.containsPoint(point):
|
|
self.setHiding()
|
|
if shape not in self.selectedShapes:
|
|
if multiple_selection_mode:
|
|
self.selectionChanged.emit(self.selectedShapes + [shape])
|
|
else:
|
|
self.selectionChanged.emit([shape])
|
|
self.hShapeIsSelected = False
|
|
else:
|
|
self.hShapeIsSelected = True
|
|
self.calculateOffsets(point)
|
|
return
|
|
self.deSelectShape()
|
|
|
|
def calculateOffsets(self, point: QPointF) -> None:
|
|
left = self.pixmap.width() - 1
|
|
right = 0
|
|
top = self.pixmap.height() - 1
|
|
bottom = 0
|
|
for s in self.selectedShapes:
|
|
rect = s.boundingRect()
|
|
if rect.left() < left:
|
|
left = rect.left()
|
|
if rect.right() > right:
|
|
right = rect.right()
|
|
if rect.top() < top:
|
|
top = rect.top()
|
|
if rect.bottom() > bottom:
|
|
bottom = rect.bottom()
|
|
|
|
x1 = left - point.x()
|
|
y1 = top - point.y()
|
|
x2 = right - point.x()
|
|
y2 = bottom - point.y()
|
|
self.offsets = QPointF(x1, y1), QPointF(x2, y2)
|
|
|
|
def boundedMoveVertex(self, pos: QPointF) -> None:
|
|
if self.hVertex is None:
|
|
logger.warning("hVertex is None, so cannot move vertex: pos=%r", pos)
|
|
return
|
|
assert self.hShape is not None
|
|
|
|
point: QPointF = self.hShape[self.hVertex]
|
|
if self.outOfPixmap(pos):
|
|
pos = self.intersectionPoint(point, pos)
|
|
self.hShape.moveVertexBy(i=self.hVertex, offset=pos - point)
|
|
|
|
def boundedMoveShapes(self, shapes, pos):
|
|
if self.outOfPixmap(pos):
|
|
return False # No need to move
|
|
o1 = pos + self.offsets[0]
|
|
if self.outOfPixmap(o1):
|
|
pos -= QPointF(min(0, o1.x()), min(0, o1.y()))
|
|
o2 = pos + self.offsets[1]
|
|
if self.outOfPixmap(o2):
|
|
pos += QPointF(
|
|
min(0, self.pixmap.width() - o2.x()),
|
|
min(0, self.pixmap.height() - o2.y()),
|
|
)
|
|
# XXX: The next line tracks the new position of the cursor
|
|
# relative to the shape, but also results in making it
|
|
# a bit "shaky" when nearing the border and allows it to
|
|
# go outside of the shape's area for some reason.
|
|
# self.calculateOffsets(self.selectedShapes, pos)
|
|
dp = pos - self.prevPoint
|
|
if dp:
|
|
for shape in shapes:
|
|
shape.moveBy(dp)
|
|
self.prevPoint = pos
|
|
return True
|
|
return False
|
|
|
|
def deSelectShape(self):
|
|
if self.selectedShapes:
|
|
self.setHiding(False)
|
|
self.selectionChanged.emit([])
|
|
self.hShapeIsSelected = False
|
|
self.update()
|
|
|
|
def deleteSelected(self):
|
|
deleted_shapes = []
|
|
if self.selectedShapes:
|
|
for shape in self.selectedShapes:
|
|
self.shapes.remove(shape)
|
|
deleted_shapes.append(shape)
|
|
self.storeShapes()
|
|
self.selectedShapes = []
|
|
self.update()
|
|
return deleted_shapes
|
|
|
|
def deleteShape(self, shape):
|
|
if shape in self.selectedShapes:
|
|
self.selectedShapes.remove(shape)
|
|
if shape in self.shapes:
|
|
self.shapes.remove(shape)
|
|
self.storeShapes()
|
|
self.update()
|
|
|
|
def paintEvent(self, a0: QtGui.QPaintEvent) -> None:
|
|
if not self.pixmap:
|
|
return super().paintEvent(a0)
|
|
|
|
p = self._painter
|
|
p.begin(self)
|
|
p.setRenderHint(QtGui.QPainter.Antialiasing)
|
|
p.setRenderHint(QtGui.QPainter.HighQualityAntialiasing)
|
|
p.setRenderHint(QtGui.QPainter.SmoothPixmapTransform)
|
|
|
|
p.scale(self.scale, self.scale)
|
|
p.translate(self.offsetToCenter())
|
|
|
|
p.drawPixmap(0, 0, self.pixmap)
|
|
|
|
p.scale(1 / self.scale, 1 / self.scale)
|
|
|
|
# draw crosshair
|
|
if (
|
|
self._crosshair[self._createMode]
|
|
and self.drawing()
|
|
and self.prevMovePoint is not None
|
|
and not self.outOfPixmap(self.prevMovePoint)
|
|
):
|
|
p.setPen(QtGui.QColor(0, 0, 0))
|
|
p.drawLine(
|
|
0,
|
|
int(self.prevMovePoint.y() * self.scale),
|
|
self.width() - 1,
|
|
int(self.prevMovePoint.y() * self.scale),
|
|
)
|
|
p.drawLine(
|
|
int(self.prevMovePoint.x() * self.scale),
|
|
0,
|
|
int(self.prevMovePoint.x() * self.scale),
|
|
self.height() - 1,
|
|
)
|
|
|
|
Shape.scale = self.scale
|
|
for shape in self.shapes:
|
|
if (shape.selected or not self._hideBackround) and self.isVisible(shape):
|
|
shape.fill = shape.selected or shape == self.hShape
|
|
shape.paint(p)
|
|
if self.current:
|
|
self.current.paint(p)
|
|
assert len(self.line.points) == len(self.line.point_labels)
|
|
self.line.paint(p)
|
|
if self.selectedShapesCopy:
|
|
for s in self.selectedShapesCopy:
|
|
s.paint(p)
|
|
|
|
if not self.current or self.createMode not in [
|
|
"polygon",
|
|
"ai_polygon",
|
|
"ai_mask",
|
|
]:
|
|
p.end()
|
|
return
|
|
|
|
drawing_shape: Shape = self.current.copy()
|
|
if self.createMode == "polygon":
|
|
if self.fillDrawing() and len(self.current.points) >= 2:
|
|
assert drawing_shape.fill_color is not None
|
|
if drawing_shape.fill_color.getRgb()[3] == 0:
|
|
logger.warning(
|
|
"fill_drawing=true, but fill_color is transparent,"
|
|
" so forcing to be opaque."
|
|
)
|
|
drawing_shape.fill_color.setAlpha(64)
|
|
drawing_shape.addPoint(self.line[1])
|
|
elif self.createMode in ["ai_polygon", "ai_mask"]:
|
|
drawing_shape.addPoint(
|
|
point=self.line.points[1],
|
|
label=self.line.point_labels[1],
|
|
)
|
|
_update_shape_with_sam(
|
|
sam=self._get_ai_model(),
|
|
image_embedding=self._get_ai_image_embedding(),
|
|
shape=drawing_shape,
|
|
createMode=self.createMode,
|
|
)
|
|
drawing_shape.fill = self.fillDrawing()
|
|
drawing_shape.selected = self.fillDrawing()
|
|
drawing_shape.paint(p)
|
|
p.end()
|
|
|
|
def transformPos(self, point: QPointF) -> QPointF:
|
|
"""Convert from widget-logical coordinates to painter-logical ones."""
|
|
return point / self.scale - self.offsetToCenter()
|
|
|
|
def enableDragging(self, enabled: bool):
|
|
self._is_dragging_enabled = enabled
|
|
|
|
def offsetToCenter(self) -> QPointF:
|
|
s = self.scale
|
|
area = super().size()
|
|
w, h = self.pixmap.width() * s, self.pixmap.height() * s
|
|
aw, ah = area.width(), area.height()
|
|
x = (aw - w) / (2 * s) if aw > w else 0
|
|
y = (ah - h) / (2 * s) if ah > h else 0
|
|
return QPointF(x, y)
|
|
|
|
def outOfPixmap(self, p: QPointF) -> bool:
|
|
w, h = self.pixmap.width(), self.pixmap.height()
|
|
return not (0 <= p.x() <= w - 1 and 0 <= p.y() <= h - 1)
|
|
|
|
def finalise(self):
|
|
assert self.current
|
|
if self.createMode in ["ai_polygon", "ai_mask"]:
|
|
_update_shape_with_sam(
|
|
sam=self._get_ai_model(),
|
|
image_embedding=self._get_ai_image_embedding(),
|
|
shape=self.current,
|
|
createMode=self.createMode,
|
|
)
|
|
self.current.close()
|
|
|
|
self.shapes.append(self.current)
|
|
self.storeShapes()
|
|
self.current = None
|
|
self.setHiding(False)
|
|
self.newShape.emit()
|
|
self.update()
|
|
|
|
def closeEnough(self, p1, p2):
|
|
# d = distance(p1 - p2)
|
|
# m = (p1-p2).manhattanLength()
|
|
# print "d %.2f, m %d, %.2f" % (d, m, d - m)
|
|
# divide by scale to allow more precision when zoomed in
|
|
return labelme.utils.distance(p1 - p2) < (self.epsilon / self.scale)
|
|
|
|
def intersectionPoint(self, p1: QPointF, p2: QPointF) -> QPointF:
|
|
# Cycle through each image edge in clockwise fashion,
|
|
# and find the one intersecting the current line segment.
|
|
# http://paulbourke.net/geometry/lineline2d/
|
|
size = self.pixmap.size()
|
|
points = [
|
|
(0, 0),
|
|
(size.width() - 1, 0),
|
|
(size.width() - 1, size.height() - 1),
|
|
(0, size.height() - 1),
|
|
]
|
|
# x1, y1 should be in the pixmap, x2, y2 should be out of the pixmap
|
|
x1 = min(max(p1.x(), 0), size.width() - 1)
|
|
y1 = min(max(p1.y(), 0), size.height() - 1)
|
|
x2, y2 = p2.x(), p2.y()
|
|
d, i, (x, y) = min(self.intersectingEdges((x1, y1), (x2, y2), points))
|
|
x3, y3 = points[i]
|
|
x4, y4 = points[(i + 1) % 4]
|
|
if (x, y) == (x1, y1):
|
|
# Handle cases where previous point is on one of the edges.
|
|
if x3 == x4:
|
|
return QPointF(x3, min(max(0, y2), max(y3, y4)))
|
|
else: # y3 == y4
|
|
return QPointF(min(max(0, x2), max(x3, x4)), y3)
|
|
return QPointF(x, y)
|
|
|
|
def intersectingEdges(self, point1, point2, points):
|
|
"""Find intersecting edges.
|
|
|
|
For each edge formed by `points', yield the intersection
|
|
with the line segment `(x1,y1) - (x2,y2)`, if it exists.
|
|
Also return the distance of `(x2,y2)' to the middle of the
|
|
edge along with its index, so that the one closest can be chosen.
|
|
"""
|
|
(x1, y1) = point1
|
|
(x2, y2) = point2
|
|
for i in range(4):
|
|
x3, y3 = points[i]
|
|
x4, y4 = points[(i + 1) % 4]
|
|
denom = (y4 - y3) * (x2 - x1) - (x4 - x3) * (y2 - y1)
|
|
nua = (x4 - x3) * (y1 - y3) - (y4 - y3) * (x1 - x3)
|
|
nub = (x2 - x1) * (y1 - y3) - (y2 - y1) * (x1 - x3)
|
|
if denom == 0:
|
|
# This covers two cases:
|
|
# nua == nub == 0: Coincident
|
|
# otherwise: Parallel
|
|
continue
|
|
ua, ub = nua / denom, nub / denom
|
|
if 0 <= ua <= 1 and 0 <= ub <= 1:
|
|
x = x1 + ua * (x2 - x1)
|
|
y = y1 + ua * (y2 - y1)
|
|
m = QPointF((x3 + x4) / 2, (y3 + y4) / 2)
|
|
d = labelme.utils.distance(m - QPointF(x2, y2))
|
|
yield d, i, (x, y)
|
|
|
|
# These two, along with a call to adjustSize are required for the
|
|
# scroll area.
|
|
def sizeHint(self):
|
|
return self.minimumSizeHint()
|
|
|
|
def minimumSizeHint(self):
|
|
if not self.pixmap:
|
|
return super().minimumSizeHint()
|
|
|
|
min_size = self.scale * self.pixmap.size()
|
|
if self._is_dragging_enabled:
|
|
# When drag buffer should be enabled, add a bit of buffer around the image
|
|
# This lets dragging the image around have a bit of give on the edges
|
|
min_size = 1.167 * min_size
|
|
return min_size
|
|
|
|
def wheelEvent(self, a0: QtGui.QWheelEvent) -> None:
|
|
mods: Qt.KeyboardModifiers = a0.modifiers()
|
|
delta: QPoint = a0.angleDelta()
|
|
if Qt.ControlModifier == int(mods):
|
|
# with Ctrl/Command key
|
|
# zoom
|
|
self.zoomRequest.emit(delta.y(), a0.posF())
|
|
else:
|
|
# scroll
|
|
self.scrollRequest.emit(delta.x(), Qt.Horizontal)
|
|
self.scrollRequest.emit(delta.y(), Qt.Vertical)
|
|
a0.accept()
|
|
|
|
def moveByKeyboard(self, offset):
|
|
if self.selectedShapes:
|
|
self.boundedMoveShapes(self.selectedShapes, self.prevPoint + offset)
|
|
self.repaint()
|
|
self.movingShape = True
|
|
|
|
def keyPressEvent(self, a0: QtGui.QKeyEvent) -> None:
|
|
modifiers = a0.modifiers()
|
|
key = a0.key()
|
|
if self.drawing():
|
|
if key == Qt.Key_Escape and self.current:
|
|
self.current = None
|
|
self.drawingPolygon.emit(False)
|
|
self.update()
|
|
elif (
|
|
key in (QtCore.Qt.Key_Return, QtCore.Qt.Key_Space)
|
|
and self.canCloseShape()
|
|
):
|
|
self.finalise()
|
|
elif modifiers == Qt.AltModifier:
|
|
self.snapping = False
|
|
elif self.editing():
|
|
if key == Qt.Key_Up:
|
|
self.moveByKeyboard(QPointF(0.0, -MOVE_SPEED))
|
|
elif key == Qt.Key_Down:
|
|
self.moveByKeyboard(QPointF(0.0, MOVE_SPEED))
|
|
elif key == Qt.Key_Left:
|
|
self.moveByKeyboard(QPointF(-MOVE_SPEED, 0.0))
|
|
elif key == Qt.Key_Right:
|
|
self.moveByKeyboard(QPointF(MOVE_SPEED, 0.0))
|
|
self._update_status()
|
|
|
|
def keyReleaseEvent(self, a0: QtGui.QKeyEvent) -> None:
|
|
modifiers = a0.modifiers()
|
|
if self.drawing():
|
|
if int(modifiers) == 0:
|
|
self.snapping = True
|
|
elif self.editing():
|
|
if self.movingShape and self.selectedShapes:
|
|
index = self.shapes.index(self.selectedShapes[0])
|
|
if self.shapesBackups[-1][index].points != self.shapes[index].points:
|
|
self.storeShapes()
|
|
self.shapeMoved.emit()
|
|
|
|
self.movingShape = False
|
|
|
|
def setLastLabel(self, text, flags):
|
|
assert text
|
|
self.shapes[-1].label = text
|
|
self.shapes[-1].flags = flags
|
|
self.shapesBackups.pop()
|
|
self.storeShapes()
|
|
return self.shapes[-1]
|
|
|
|
def undoLastLine(self):
|
|
assert self.shapes
|
|
self.current = self.shapes.pop()
|
|
self.current.setOpen()
|
|
self.current.restoreShapeRaw()
|
|
if self.createMode in ["polygon", "linestrip"]:
|
|
self.line.points = [self.current[-1], self.current[0]]
|
|
elif self.createMode in ["rectangle", "line", "circle"]:
|
|
self.current.points = self.current.points[0:1]
|
|
elif self.createMode == "point":
|
|
self.current = None
|
|
self.drawingPolygon.emit(True)
|
|
|
|
def undoLastPoint(self):
|
|
if not self.current or self.current.isClosed():
|
|
return
|
|
self.current.popPoint()
|
|
if len(self.current) > 0:
|
|
self.line[0] = self.current[-1]
|
|
else:
|
|
self.current = None
|
|
self.drawingPolygon.emit(False)
|
|
self.update()
|
|
|
|
def loadPixmap(self, pixmap, clear_shapes=True):
|
|
self.pixmap = pixmap
|
|
if clear_shapes:
|
|
self.shapes = []
|
|
self.update()
|
|
|
|
def loadShapes(self, shapes, replace=True):
|
|
if replace:
|
|
self.shapes = list(shapes)
|
|
else:
|
|
self.shapes.extend(shapes)
|
|
self.storeShapes()
|
|
self.current = None
|
|
self.hShape = None
|
|
self.hVertex = None
|
|
self.hEdge = None
|
|
self.update()
|
|
|
|
def setShapeVisible(self, shape, value):
|
|
self.visible[shape] = value
|
|
self.update()
|
|
|
|
def overrideCursor(self, cursor):
|
|
if cursor == self._cursor:
|
|
return
|
|
self.restoreCursor()
|
|
self._cursor = cursor
|
|
QtWidgets.QApplication.setOverrideCursor(cursor)
|
|
|
|
def restoreCursor(self):
|
|
self._cursor = CURSOR_DEFAULT
|
|
QtWidgets.QApplication.restoreOverrideCursor()
|
|
|
|
def resetState(self):
|
|
self.restoreCursor()
|
|
self.pixmap = QtGui.QPixmap()
|
|
self.shapes = []
|
|
self.shapesBackups = []
|
|
self.movingShape = False
|
|
self.selectedShapes = []
|
|
self.selectedShapesCopy = []
|
|
self.current = None
|
|
self.hShape = None
|
|
self.prevhShape = None
|
|
self.hVertex = None
|
|
self.prevhVertex = None
|
|
self.hEdge = None
|
|
self.prevhEdge = None
|
|
self.update()
|
|
|
|
|
|
def _update_shape_with_sam(
|
|
sam: osam.types.Model,
|
|
image_embedding: osam.types.ImageEmbedding,
|
|
shape: Shape,
|
|
createMode: Literal["ai_polygon", "ai_mask"],
|
|
) -> None:
|
|
if createMode not in ["ai_polygon", "ai_mask"]:
|
|
raise ValueError(
|
|
f"createMode must be 'ai_polygon' or 'ai_mask', not {createMode}"
|
|
)
|
|
|
|
response: osam.types.GenerateResponse = sam.generate(
|
|
request=osam.types.GenerateRequest(
|
|
model=sam.name,
|
|
image_embedding=image_embedding,
|
|
prompt=osam.types.Prompt(
|
|
points=np.array([[point.x(), point.y()] for point in shape.points]),
|
|
point_labels=np.array(shape.point_labels),
|
|
),
|
|
)
|
|
)
|
|
if not response.annotations:
|
|
logger.warning("No annotations returned by model {!r}", sam)
|
|
return
|
|
|
|
if createMode == "ai_mask":
|
|
y1: int
|
|
x1: int
|
|
y2: int
|
|
x2: int
|
|
if response.annotations[0].bounding_box is None:
|
|
y1, x1, y2, x2 = imgviz.instances.mask_to_bbox(
|
|
[response.annotations[0].mask]
|
|
)[0].astype(int)
|
|
else:
|
|
y1 = response.annotations[0].bounding_box.ymin
|
|
x1 = response.annotations[0].bounding_box.xmin
|
|
y2 = response.annotations[0].bounding_box.ymax
|
|
x2 = response.annotations[0].bounding_box.xmax
|
|
shape.setShapeRefined(
|
|
shape_type="mask",
|
|
points=[QPointF(x1, y1), QPointF(x2, y2)],
|
|
point_labels=[1, 1],
|
|
mask=response.annotations[0].mask[y1 : y2 + 1, x1 : x2 + 1],
|
|
)
|
|
elif createMode == "ai_polygon":
|
|
points = polygon_from_mask.compute_polygon_from_mask(
|
|
mask=response.annotations[0].mask
|
|
)
|
|
if len(points) < 2:
|
|
return
|
|
shape.setShapeRefined(
|
|
shape_type="polygon",
|
|
points=[QPointF(point[0], point[1]) for point in points],
|
|
point_labels=[1] * len(points),
|
|
) |