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