import imgviz from qtpy import QtCore from qtpy import QtGui from qtpy import QtWidgets import labelme.ai import labelme.utils from labelme import QT5 from labelme.logger import logger from labelme.shape import Shape import numpy as np import cv2 # TODO(unknown): # - [maybe] Find optimal epsilon value. CURSOR_DEFAULT = QtCore.Qt.ArrowCursor CURSOR_POINT = QtCore.Qt.PointingHandCursor CURSOR_DRAW = QtCore.Qt.CrossCursor CURSOR_MOVE = QtCore.Qt.ClosedHandCursor CURSOR_GRAB = QtCore.Qt.OpenHandCursor MOVE_SPEED = 5.0 class Canvas(QtWidgets.QWidget): zoomRequest = QtCore.Signal(int, QtCore.QPoint) scrollRequest = QtCore.Signal(int, int) newShape = QtCore.Signal() selectionChanged = QtCore.Signal(list) shapeMoved = QtCore.Signal() drawingPolygon = QtCore.Signal(bool) vertexSelected = QtCore.Signal(bool) mouseMoved = QtCore.Signal(QtCore.QPointF) CREATE, EDIT = 0, 1 # polygon, rectangle, line, or point _createMode = "polygon" _fill_drawing = False def __init__(self, *args, **kwargs): self.epsilon = kwargs.pop("epsilon", 10.0) self.detection_model_path=kwargs.pop("detection_model_path",None) self.segmentation_model_path=kwargs.pop("segmentation_model_path",None) self.double_click = kwargs.pop("double_click", "close") if self.double_click not in [None, "close"]: raise ValueError( "Unexpected value for double_click event: {}".format(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(Canvas, self).__init__(*args, **kwargs) # Initialise local state. self.mode = self.EDIT self.shapes = [] self.shapesBackups = [] self.current = None self.selectedShapes = [] # save the selected shapes here self.selectedShapesCopy = [] # 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 = QtCore.QPoint() self.prevMovePoint = QtCore.QPoint() self.offsets = QtCore.QPoint(), QtCore.QPoint() self.scale = 1.0 self.pixmap = QtGui.QPixmap() self.visible = {} self._hideBackround = False self.hideBackround = False self.hShape = None self.prevhShape = None self.hVertex = None self.prevhVertex = None self.hEdge = None self.prevhEdge = None self.movingShape = False self.snapping = True self.hShapeIsSelected = False self._painter = QtGui.QPainter() self._cursor = CURSOR_DEFAULT self.draw_pred=False self.pred_bbox_points=None self.current_bbox_point=None # 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(QtCore.Qt.WheelFocus) self._ai_model = None self._detection_model = None self._segmentation_model = None 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("Unsupported createMode: %s" % value) self._createMode = value def initializeBarcodeModel(self, detection_model_path, segmentation_model_path=None): if not detection_model_path: raise ValueError("Detection model path is required.") logger.debug("Initializing only detection model: %r" % "BarcodePredictModel") self._detection_model = labelme.ai.BarcodePredictModel(detection_model_path) if segmentation_model_path: logger.debug("Initializing barcode detection & Segmentation model: %r" % "BarcodePredictModel") self._segmentation_model = labelme.ai.BarcodePredictModel( detection_model_path, segmentation_model_path ) if self.pixmap is None: logger.warning("Pixmap is not set yet") return self._detection_model.set_image( image=labelme.utils.img_qt_to_arr(self.pixmap.toImage()) ) def initializeAiModel(self, name, weight_path=None): if self._ai_model is not None: logger.debug("AI model is already initialized.") return if name not in [model.name for model in labelme.ai.MODELS]: raise ValueError("Unsupported AI model: %s" % name) model_class = [model for model in labelme.ai.MODELS if model.name == name][0] logger.debug(f"Initializing AI model: {name}") self._ai_model = model_class(weight_path) if self.pixmap is None: logger.warning("Pixmap is not set yet") return self._ai_model.set_image( image=labelme.utils.img_qt_to_arr(self.pixmap.toImage()) # image=self.pixmap.toImage() ) 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 print(f"shape is restorable") 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, ev): self.overrideCursor(self._cursor) def leaveEvent(self, ev): self.unHighlight() self.restoreCursor() def focusOutEvent(self, ev): self.restoreCursor() def isVisible(self, shape): return self.visible.get(shape, True) def drawing(self): return self.mode == self.CREATE def editing(self): return self.mode == self.EDIT def setEditing(self, value=True): self.mode = self.EDIT if value else self.CREATE if self.mode == self.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 mouseMoveEvent(self, ev): """Update line with last point and current coordinates.""" try: if QT5: pos = self.transformPos(ev.localPos()) else: pos = self.transformPos(ev.posF()) except AttributeError: return self.mouseMoved.emit(pos) self.prevMovePoint = pos self.restoreCursor() is_shift_pressed = ev.modifiers() & QtCore.Qt.ShiftModifier # 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 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() return # Polygon copy moving. if QtCore.Qt.RightButton & ev.buttons(): if self.selectedShapesCopy and self.prevPoint: 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() return # Polygon/Vertex moving. if QtCore.Qt.LeftButton & ev.buttons(): if self.selectedVertex(): self.boundedMoveVertex(pos) self.repaint() self.movingShape = True elif self.selectedShapes and self.prevPoint: 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. self.setToolTip(self.tr("Image")) for shape in reversed([s for s in self.shapes if self.isVisible(s)]): # 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 / self.scale) index_edge = shape.nearestEdge(pos, self.epsilon / self.scale) if index is not None: if self.selectedVertex(): 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) self.setToolTip(self.tr("Click & drag to move point")) self.setStatusTip(self.toolTip()) self.update() break elif index_edge is not None and shape.canAddPoint(): if self.selectedVertex(): 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) self.setToolTip(self.tr("Click to create point")) self.setStatusTip(self.toolTip()) self.update() break elif len(shape.points)!=0 and shape.containsPoint(pos) : if self.selectedVertex(): self.hShape.highlightClear() self.prevhVertex = self.hVertex self.hVertex = None self.prevhShape = self.hShape = shape self.prevhEdge = self.hEdge self.hEdge = None self.setToolTip( self.tr("Click & drag to move shape '%s'") % shape.label ) self.setStatusTip(self.toolTip()) self.overrideCursor(CURSOR_GRAB) self.update() break else: # Nothing found, clear highlights, reset state. self.unHighlight() self.vertexSelected.emit(self.hVertex is not None) 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, ev): if QT5: pos = self.transformPos(ev.localPos()) else: pos = self.transformPos(ev.posF()) is_shift_pressed = ev.modifiers() & QtCore.Qt.ShiftModifier if ev.button() == QtCore.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(ev.modifiers()) == QtCore.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 ev.modifiers() & QtCore.Qt.ControlModifier: self.finalise() elif not self.outOfPixmap(pos): # 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 ev.modifiers() & QtCore.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(): self.addPointToEdge() elif ( self.selectedVertex() and int(ev.modifiers()) == QtCore.Qt.ShiftModifier ): # Delete point if: left-click + SHIFT on a point self.removeSelectedPoint() group_mode = int(ev.modifiers()) == QtCore.Qt.ControlModifier self.selectShapePoint(pos, multiple_selection_mode=group_mode) self.prevPoint = pos self.repaint() elif ev.button() == QtCore.Qt.RightButton and self.editing(): group_mode = int(ev.modifiers()) == QtCore.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 def mouseReleaseEvent(self, ev): if ev.button() == QtCore.Qt.RightButton: menu = self.menus[len(self.selectedShapesCopy) > 0] self.restoreCursor() if not menu.exec_(self.mapToGlobal(ev.pos())) and self.selectedShapesCopy: # Cancel the move by deleting the shadow copy. self.selectedShapesCopy = [] self.repaint() elif ev.button() == QtCore.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] ) 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 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): return self.drawing() and ( (self.current and len(self.current) > 2) or self.createMode in ["ai_polygon", "ai_mask"] ) def mouseDoubleClickEvent(self, ev): if self.double_click != "close": return if ( self.createMode == "polygon" and self.canCloseShape() ) or self.createMode in ["ai_polygon", "ai_mask"]: 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.selectedVertex(): # A vertex is marked for selection. index, shape = self.hVertex, self.hShape shape.highlightVertex(index, shape.MOVE_VERTEX) else: 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): 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 = QtCore.QPointF(x1, y1), QtCore.QPointF(x2, y2) def boundedMoveVertex(self, pos): index, shape = self.hVertex, self.hShape point = shape[index] if self.outOfPixmap(pos): pos = self.intersectionPoint(point, pos) shape.moveVertexBy(index, 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 -= QtCore.QPointF(min(0, o1.x()), min(0, o1.y())) o2 = pos + self.offsets[1] if self.outOfPixmap(o2): pos += QtCore.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 duplicateSelectedShapes(self): if self.selectedShapes: self.selectedShapesCopy = [s.copy() for s in self.selectedShapes] self.boundedShiftShapes(self.selectedShapesCopy) self.endMove(copy=True) return self.selectedShapes def boundedShiftShapes(self, shapes): # Try to move in one direction, and if it fails in another. # Give up if both fail. point = shapes[0][0] offset = QtCore.QPointF(2.0, 2.0) self.offsets = QtCore.QPoint(), QtCore.QPoint() self.prevPoint = point if not self.boundedMoveShapes(shapes, point - offset): self.boundedMoveShapes(shapes, point + offset) def paintEvent(self, event): if not self.pixmap: return super(Canvas, self).paintEvent(event) 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) # draw crosshair if ( self._crosshair[self._createMode] and self.drawing() and self.prevMovePoint and not self.outOfPixmap(self.prevMovePoint) ): p.setPen(QtGui.QColor(0, 0, 0)) p.drawLine( 0, int(self.prevMovePoint.y()), self.width() - 1, int(self.prevMovePoint.y()), ) p.drawLine( int(self.prevMovePoint.x()), 0, int(self.prevMovePoint.x()), 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(self.draw_pred and self.current is not None): # print("pred mode on") # for bbox_points in self.pred_bbox_points: # drawing_shape = self.current.copy() # drawing_shape.setShapeRefined( # shape_type="polygon", # points=[QtCore.QPointF(point[0], point[1]) for point in bbox_points], # point_labels=[1]*len(bbox_points) # ) # drawing_shape.fill = self.fillDrawing() # drawing_shape.selected = True # drawing_shape.paint(p) # self.draw_pred=False if ( self.fillDrawing() and self.createMode == "polygon" and self.current is not None and len(self.current.points) >= 2 ): drawing_shape = self.current.copy() 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]) drawing_shape.fill = True drawing_shape.paint(p) elif self.createMode == "ai_polygon" and self.current is not None: drawing_shape = self.current.copy() drawing_shape.addPoint( point=self.line.points[1], label=self.line.point_labels[1], ) points = self._detection_model.predict_polygon_from_points( points=[[point.x(), point.y()] for point in drawing_shape.points], point_labels=drawing_shape.point_labels, ) if len(points) > 2: drawing_shape.setShapeRefined( shape_type="polygon", points=[QtCore.QPointF(point[0], point[1]) for point in points], point_labels=[1] * len(points), ) drawing_shape.fill = self.fillDrawing() drawing_shape.selected = True drawing_shape.paint(p) elif self.createMode == "ai_mask" and self.current is not None: drawing_shape = self.current.copy() drawing_shape.addPoint( point=self.line.points[1], label=self.line.point_labels[1], ) mask = self._detection_model.predict_mask_from_points( points=[[point.x(), point.y()] for point in drawing_shape.points], point_labels=drawing_shape.point_labels, ) y1, x1, y2, x2 = imgviz.instances.masks_to_bboxes([mask])[0].astype(int) drawing_shape.setShapeRefined( shape_type="mask", points=[QtCore.QPointF(x1, y1), QtCore.QPointF(x2, y2)], point_labels=[1, 1], mask=mask[y1 : y2 + 1, x1 : x2 + 1], ) drawing_shape.selected = True drawing_shape.paint(p) p.end() def transformPos(self, point): """Convert from widget-logical coordinates to painter-logical ones.""" return point / self.scale - self.offsetToCenter() def offsetToCenter(self): s = self.scale area = super(Canvas, self).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 QtCore.QPointF(x, y) def outOfPixmap(self, p): w, h = self.pixmap.width(), self.pixmap.height() return not (0 <= p.x() <= w - 1 and 0 <= p.y() <= h - 1) def finalise(self): if(self.current is None): return if self.createMode == "ai_polygon": # convert points to polygon by an AI model assert self.current.shape_type == "points" points = self._detection_model.predict_polygon_from_points( points=[[point.x(), point.y()] for point in self.current.points], point_labels=self.current.point_labels, ) self.current.setShapeRefined( points=[QtCore.QPointF(point[0], point[1]) for point in points], point_labels=[1] * len(points), shape_type="polygon", ) elif self.createMode == "ai_mask": # convert points to mask by an AI model assert self.current.shape_type == "points" mask = self._detection_model.predict_mask_from_points( points=[[point.x(), point.y()] for point in self.current.points], point_labels=self.current.point_labels, ) y1, x1, y2, x2 = imgviz.instances.masks_to_bboxes([mask])[0].astype(int) self.current.setShapeRefined( shape_type="mask", points=[QtCore.QPointF(x1, y1), QtCore.QPointF(x2, y2)], point_labels=[1, 1], mask=mask[y1 : y2 + 1, x1 : x2 + 1], ) elif self.pred_bbox_points is not None and self.draw_pred: print("pred mode on") current_copy=self.current.copy() for bbox_point in self.pred_bbox_points: drawing_shape=current_copy.copy() drawing_shape.setShapeRefined( shape_type="polygon", points=[QtCore.QPointF(point[0], point[1]) for point in bbox_point], point_labels=[1]*len(bbox_point) ) drawing_shape.close() self.shapes.append(drawing_shape) self.storeShapes() self.update() self.newShape.emit() current_copy.close() current_copy=None if(self.current): self.current.close() self.current = None self.setHiding(False) self.draw_pred=False return 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, p2): # 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 QtCore.QPointF(x3, min(max(0, y2), max(y3, y4))) else: # y3 == y4 return QtCore.QPointF(min(max(0, x2), max(x3, x4)), y3) return QtCore.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 = QtCore.QPointF((x3 + x4) / 2, (y3 + y4) / 2) d = labelme.utils.distance(m - QtCore.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 self.pixmap: return self.scale * self.pixmap.size() return super(Canvas, self).minimumSizeHint() def wheelEvent(self, ev): if QT5: mods = ev.modifiers() delta = ev.angleDelta() if QtCore.Qt.ControlModifier == int(mods): # with Ctrl/Command key # zoom self.zoomRequest.emit(delta.y(), ev.pos()) else: # scroll self.scrollRequest.emit(delta.x(), QtCore.Qt.Horizontal) self.scrollRequest.emit(delta.y(), QtCore.Qt.Vertical) else: if ev.orientation() == QtCore.Qt.Vertical: mods = ev.modifiers() if QtCore.Qt.ControlModifier == int(mods): # with Ctrl/Command key self.zoomRequest.emit(ev.delta(), ev.pos()) else: self.scrollRequest.emit( ev.delta(), QtCore.Qt.Horizontal if (QtCore.Qt.ShiftModifier == int(mods)) else QtCore.Qt.Vertical, ) else: self.scrollRequest.emit(ev.delta(), QtCore.Qt.Horizontal) ev.accept() def moveByKeyboard(self, offset): if self.selectedShapes: self.boundedMoveShapes(self.selectedShapes, self.prevPoint + offset) self.repaint() self.movingShape = True def keyPressEvent(self, ev): modifiers = ev.modifiers() key = ev.key() if self.drawing(): if key == QtCore.Qt.Key_Escape and self.current: self.current = None self.drawingPolygon.emit(False) self.update() elif key == QtCore.Qt.Key_Return and self.canCloseShape(): self.finalise() elif modifiers == QtCore.Qt.AltModifier: self.snapping = False elif key == QtCore.Qt.Key_V: if self._detection_model is None: logger.info(f"Initializing AI model") self.initializeBarcodeModel(self.detection_model_path, self.segmentation_model_path) self.current = Shape( shape_type="points" if self.createMode in ["ai_polygon", "ai_mask"] else self.createMode ) if self._detection_model: if self._segmentation_model is None: logger.info(f"Performing detection only.") # Get prediction from model self.pred_bbox_points = self._detection_model.predict_polygon_from_points() print("Predicted Bounding Box Points:", self.pred_bbox_points) if self.pred_bbox_points: self.draw_pred = True self.finalise() else: print("No bounding boxes detected.") else: logger.info(f"Performing detection and segmentation.") self.detect_and_segment() elif self.editing(): if key == QtCore.Qt.Key_Up: self.moveByKeyboard(QtCore.QPointF(0.0, -MOVE_SPEED)) elif key == QtCore.Qt.Key_Down: self.moveByKeyboard(QtCore.QPointF(0.0, MOVE_SPEED)) elif key == QtCore.Qt.Key_Left: self.moveByKeyboard(QtCore.QPointF(-MOVE_SPEED, 0.0)) elif key == QtCore.Qt.Key_Right: self.moveByKeyboard(QtCore.QPointF(MOVE_SPEED, 0.0)) def scale_points(self, approx, mask_shape, cropped_shape, x_min, y_min): scale_x = cropped_shape[1] / mask_shape[1] # Scale factor for x-axis scale_y = cropped_shape[0] / mask_shape[0] # Scale factor for y-axis return [[int(pt[0][0] * scale_x) + x_min, int(pt[0][1] * scale_y) + y_min] for pt in approx] def detect_and_segment(self): """ Perform detection and segmentation (if both models are available). """ logger.info("Performing detection and segmentation.") self.current = Shape( shape_type="points" if self.createMode in ["ai_polygon", "ai_mask"] else self.createMode ) # Step 1: detection bounding box points detection_results = self._detection_model.predict_polygon_from_points() if not detection_results or len(detection_results) == 0: logger.warning("No detection found") return logger.debug(f"Detection results: {detection_results}") all_segmentation_results = [] rotated = False # Step 2: Loop through each detection result since there are multiple per image for detection_idx, detection_result in enumerate(detection_results): logger.debug(f"Processing detection {detection_idx + 1}/{len(detection_results)}") try: #extracting x_coords = [point[0] for point in detection_result] y_coords = [point[1] for point in detection_result] #min and max values for x and y x_min, x_max = min(x_coords), max(x_coords) y_min, y_max = min(y_coords), max(y_coords) logger.debug(f"Bounding box for detection {detection_idx + 1} - x_min: {x_min}, y_min: {y_min}, x_max: {x_max}, y_max: {y_max}") except Exception as e: logger.error(f"Error extracting bounding box coordinates for detection {detection_idx + 1}: {e}") continue # Converting bounding box values to integers for cropping x_min, y_min, x_max, y_max = int(x_min), int(y_min), int(x_max), int(y_max) # Step 3: Cropping image based on detection output try: cropped_image = self.pixmap.toImage().copy(x_min, y_min, x_max - x_min, y_max - y_min) cropped_image = labelme.utils.img_qt_to_arr(cropped_image) orig_height, orig_width = cropped_image.shape[:2] logger.debug(f"Original height: {orig_height}, Original width: {orig_width}") # if the height is greater than the width we rotate for segmentaion if orig_height > orig_width: cropped_image = cv2.rotate(cropped_image, cv2.ROTATE_90_CLOCKWISE) logger.debug(f"Rotated cropped image for detection {detection_idx + 1} due to height > width.") orig_cropped_shape = cropped_image.shape[:2] rotated = True else: rotated = False # Save crop image cv2.imwrite(f"cropped_image_{detection_idx + 1}.png", cropped_image) logger.debug(f"Saved cropped image for detection {detection_idx + 1}: {cropped_image.shape}") # logger.debug(f"Cropped image shape for detection {detection_idx + 1}: {cropped_image.shape}") except Exception as e: logger.error(f"Error cropping the image for detection {detection_idx + 1}: {e}") continue # Step 4: Resize the cropped image to match segmentation input size (1 64 256) try: orig_cropped_shape = cropped_image.shape[:2] # Save the original cropped image size preprocessed_img = self._detection_model.preprocess_image(cropped_image, for_segmentation=True) logger.debug(f"Preprocessed image shape for segmentation detection {detection_idx + 1}: {preprocessed_img.shape}") except Exception as e: logger.error(f"Error preprocessing the image for segmentation for detection {detection_idx + 1}: {e}") continue # Step 5: inference on segmentation model on cropped image try: seg_result = self._segmentation_model.segmentation_sess.infer_new_request({'x': preprocessed_img}) logger.debug(f"Segmentation model inference completed for detection {detection_idx + 1}.") except Exception as e: logger.error(f"Error during segmentation model inference for detection {detection_idx + 1}: {e}") continue # Step 6: Convert binary mask to polygon (contours) try: mask = seg_result['save_infer_model/scale_0.tmp_0'] #model output name mask = mask.squeeze() # Remove batch dimension, should result in (64, 256) logger.debug(f"Segmentation mask shape for detection {detection_idx + 1}: {mask.shape}") # Normalize the mask to 0 and 255 and convert to uint8 mask = (mask * 255).astype(np.uint8) logger.debug(f"Converted mask shape for detection {detection_idx + 1}: {mask.shape}, dtype: {mask.dtype}") cv2.imwrite(f"segmentation_mask_{detection_idx + 1}.png", mask) if rotated: cropped_image = cv2.rotate(cropped_image, cv2.ROTATE_90_COUNTERCLOCKWISE) mask = cv2.rotate(mask, cv2.ROTATE_90_COUNTERCLOCKWISE) rotated_cropped_shape = cropped_image.shape[:2] # cv2.imwrite(f"segmentation_mask_{detection_idx + 1}.png", mask) logger.debug(f"Saved segmentation mask for detection {detection_idx + 1}.") # Step 7: Find contours contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) logger.debug(f"Found {len(contours)} contours in the mask for detection {detection_idx + 1}.") if len(contours) > 0: largest_contour = max(contours, key=cv2.contourArea) # Step 8: Approximate a polygon with exactly 4 points (quadrilateral) epsilon = 0.02 * cv2.arcLength(largest_contour, True) # epsilon for precision approx = cv2.approxPolyDP(largest_contour, epsilon, True) # If the approximation doesn't result in 4 points, force it if len(approx) != 4: # Using boundingRect as fallback in case of insufficient points print("log here") x, y, w, h = cv2.boundingRect(largest_contour) point_xy = [ [x + x_min, y + y_min], # Top-left [x + w + x_min, y + y_min], # Top-right [x + w + x_min, y + h + y_min], # Bottom-right [x + x_min, y + h + y_min] # Bottom-left ] else: if rotated: point_xy = self.scale_points(approx, mask.shape, rotated_cropped_shape, x_min, y_min) else: point_xy = self.scale_points(approx, mask.shape, orig_cropped_shape, x_min, y_min) logger.debug(f"Generated 4 corner points for the polygon for detection {detection_idx + 1}: {point_xy}") self.pred_bbox_points = [point_xy] logger.debug(f"Predicted Bounding Box Points for detection {detection_idx + 1}: {self.pred_bbox_points}") if self.pred_bbox_points: self.draw_pred = True self.finalise() else: logger.info(f"No bounding boxes detected for detection {detection_idx + 1}.") # Collect the segmentation result all_segmentation_results.append(self.pred_bbox_points) except Exception as e: logger.error(f"Error creating the polygon shape for detection {detection_idx + 1}: {e}") # **Reset critical variables after each detection**: self.pred_bbox_points = None self.draw_pred = False # You now have a list of segmentation results for all detections if all_segmentation_results: logger.info(f"Segmentation results for all detections: {all_segmentation_results}") def keyReleaseEvent(self, ev): modifiers = ev.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 if(self.shapes is None or len(self.shapes)==0): return self.shapes[-1].label = text self.shapes[-1].flags = flags self.shapesBackups.pop() self.storeShapes() return self.shapes[-1] def undoLastLine(self): if(self.shapes is None or len(self.shapes)==0): return self.current = self.shapes.pop() self.current.setOpen() self.current.restoreShapeRaw() if self.createMode in ["polygon", "linestrip"] and self.draw_pred is False: 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 self._detection_model: self._detection_model.set_image( image=labelme.utils.img_qt_to_arr(self.pixmap.toImage()) ) 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): self.restoreCursor() self._cursor = cursor QtWidgets.QApplication.setOverrideCursor(cursor) def restoreCursor(self): QtWidgets.QApplication.restoreOverrideCursor() def resetState(self): self.restoreCursor() self.pixmap = None self.shapesBackups = [] self.update()