1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405 |
- import imgviz
- from qtpy import QtCore
- from qtpy import QtGui
- from qtpy import QtWidgets
- import cv2
- import numpy as np
- 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
- from labelme.utils import load_barcode_dict
- from labelme.widgets.label_dialog import LabelDialog
- # 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.decoding_model_path=kwargs.pop("decoding_model_path", None)
- self.double_click = kwargs.pop("double_click", "close")
- self._characters = load_barcode_dict() # Load dictionary once
- logger.debug(f"Loaded barcode dictionary with {len(self._characters)} characters.")
- 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
- self.decoded_barcode = 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
- self._decoding_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, decoding_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.BarcodeDetectModel(detection_model_path)
- if segmentation_model_path:
- logger.debug("Initializing barcode detection & Segmentation model: %r" % "BarcodePredictModel")
- self._segmentation_model = labelme.ai.BarcodeDetectModel(
- detection_model_path, segmentation_model_path
- )
- if decoding_model_path:
- logger.debug("Initializing barcode detection, Segmentation model, Decoding_model: %r" % "BarcodePredictModel")
- self._decoding_model = labelme.ai.BarcodeDecodeModel(
- decoding_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()
- # print("self.line.points", self.line.points)
- 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 in manual mode and points are finalized
- if self.createMode in ["polygon", "rectangle"] and len(self.current.points) > 0:
- if self._detection_model is None:
- logger.info(f"Initializing AI model")
- self.initializeBarcodeModel(self.detection_model_path, self.segmentation_model_path, self.decoding_model_path)
- # Extract the points from the manually drawn shape
- manual_points = [[point.x(), point.y()] for point in self.current.points]
- print(manual_points)
- logger.debug(f"Manually drawn points: {manual_points}")
- self.detect_and_segment(manual_points=manual_points)
- self.mode = self.EDIT # Or any appropriate mode to disable drawing behavior
- 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.description = self.decoded_barcode
- drawing_shape.close()
- self.shapes.append(drawing_shape)
- self.storeShapes()
- self.update()
- self.newShape.emit()
- # self.pred_bbox_points = None
- # self.draw_pred = False
- # self.detect_and_segment() # Process the next detection
- 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_path 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, segmentation and decoding.")
- self.initializeBarcodeModel(self.detection_model_path, self.segmentation_model_path, self.decoding_model_path)
- self.detect_and_segment(manual_points=None)
- 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, points, mask_shape, cropped_shape, x_min, y_min):
- """
- Scale the points from the segmentation mask space back to the original cropped image coordinates.
-
- Args:
- points (np.ndarray): Points to scale (Nx2 array or list of lists).
- mask_shape (tuple): Shape of the segmentation mask (height, width).
- cropped_shape (tuple): Shape of the cropped image in the original image (height, width).
- x_min (int): Minimum x-coordinate of the cropped region in the original image.
- y_min (int): Minimum y-coordinate of the cropped region in the original image.
- Returns:
- scaled_points (list): List of scaled points as [[x1, y1], [x2, y2], ...].
- """
- # Compute scaling factors
- 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
- # Scale and translate the points
- scaled_points = []
- for pt in points:
- scaled_x = int(pt[0] * scale_x) + x_min
- scaled_y = int(pt[1] * scale_y) + y_min
- scaled_points.append([scaled_x, scaled_y])
- # logger.debug(f"scaled points: {scaled_points}")
- return scaled_points
-
- def expand_bbox(self, x_min, y_min, x_max, y_max, factor=1.3):
- """
- Expands the bounding box by a given factor.
- """
- center_x = (x_max + x_min) / 2
- center_y = (y_max + y_min) / 2
- width = (x_max - x_min) * factor
- height = (y_max - y_min) * factor
- new_x_min = int(center_x - width / 2)
- new_y_min = int(center_y - height / 2)
- new_x_max = int(center_x + width / 2)
- new_y_max = int(center_y + height / 2)
- # Ensure the coordinates do not go out of image bounds
- new_x_min = max(0, new_x_min)
- new_y_min = max(0, new_y_min)
- new_x_max = min(self.pixmap.width() - 1, new_x_max)
- new_y_max = min(self.pixmap.height() - 1, new_y_max)
- return new_x_min, new_y_min, new_x_max, new_y_max
- def detect_and_segment(self, manual_points=None):
- """
- 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
- )
- print("self.current",self.current)
- # Step 1: detection bounding box points
- detection_results = [manual_points] if manual_points else self._detection_model.predict_polygon_from_points()
- print(detection_results)
- 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 = self.expand_bbox(x_min, y_min, x_max, 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)
-
- # Scale mask to original cropped image size
- scaled_mask = cv2.resize(mask, (orig_cropped_shape[1], orig_cropped_shape[0]), interpolation=cv2.INTER_NEAREST)
-
- if rotated:
- cropped_image = cv2.rotate(cropped_image, cv2.ROTATE_90_COUNTERCLOCKWISE)
- mask = cv2.rotate(scaled_mask, cv2.ROTATE_90_COUNTERCLOCKWISE)
- rotated_cropped_shape = cropped_image.shape[:2]
- else:
- mask = scaled_mask
-
- cv2.imwrite(f"scaled_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:
- # Get the largest contour
- largest_contour = max(contours, key=cv2.contourArea)
- # Get bounding and rotated rectangles
- rot_rect = cv2.minAreaRect(largest_contour)
- box_points = cv2.boxPoints(rot_rect)
- box_points = np.int32(box_points)
- # Log the rotated rectangle points
- # logger.debug(f"Rotated corners (box_points) for detection {detection_idx + 1}: {box_points}")
- # Step 1: Draw the rotated rectangle on the cropped image
- cropped_with_rects = cropped_image.copy()
- cv2.drawContours(cropped_with_rects, [box_points], -1, (0, 255, 0), 2)
- cv2.imwrite(f"cropped_with_rects_{detection_idx + 1}.png", cropped_with_rects)
- logger.debug(f"Saved cropped image with rectangles for detection {detection_idx + 1}.")
- if rotated:
- # Scale points from mask space to the original image space with rotation
- scaled_points = self.scale_points(box_points, mask.shape, rotated_cropped_shape, x_min, y_min)
- else:
- # Scale points without rotation
- scaled_points = self.scale_points(box_points, mask.shape, orig_cropped_shape, x_min, y_min)
- self.pred_bbox_points = [scaled_points]
- logger.debug(f"Predicted Bounding Box Points for barcode {detection_idx + 1}: {self.pred_bbox_points}")
-
- if self.pred_bbox_points:
- for bbox_points in self.pred_bbox_points:
- try:
- original_image = labelme.utils.img_qt_to_arr(self.pixmap.toImage())
- # Call the decoding function
- decoded_output = self._decoding_model.decode_from_points(
- points=bbox_points,
- detection_idx=detection_idx,
- original_image = original_image
- )
- logger.debug(f"Decoded output for detection {detection_idx + 1}: {decoded_output}")
- #passed to popup through app.py and label_dailog.py
- self.decoded_barcode = decoded_output
- # (Optional) Use the decoded output for updating UI or further processing
- # Example: Print or log the decoded text
- # print(f"Decoded output: {decoded_output}")
- except Exception as e:
- logger.error(f"Error during decoding for bbox points {bbox_points}: {e}")
-
- ## pass each shape points to decode function here, do the decodign and give back the output here and print the output
- 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 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()
|