canvas.py 57 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405
  1. import imgviz
  2. from qtpy import QtCore
  3. from qtpy import QtGui
  4. from qtpy import QtWidgets
  5. import cv2
  6. import numpy as np
  7. import labelme.ai
  8. import labelme.utils
  9. from labelme import QT5
  10. from labelme.logger import logger
  11. from labelme.shape import Shape
  12. import numpy as np
  13. import cv2
  14. from labelme.utils import load_barcode_dict
  15. from labelme.widgets.label_dialog import LabelDialog
  16. # TODO(unknown):
  17. # - [maybe] Find optimal epsilon value.
  18. CURSOR_DEFAULT = QtCore.Qt.ArrowCursor
  19. CURSOR_POINT = QtCore.Qt.PointingHandCursor
  20. CURSOR_DRAW = QtCore.Qt.CrossCursor
  21. CURSOR_MOVE = QtCore.Qt.ClosedHandCursor
  22. CURSOR_GRAB = QtCore.Qt.OpenHandCursor
  23. MOVE_SPEED = 5.0
  24. class Canvas(QtWidgets.QWidget):
  25. zoomRequest = QtCore.Signal(int, QtCore.QPoint)
  26. scrollRequest = QtCore.Signal(int, int)
  27. newShape = QtCore.Signal()
  28. selectionChanged = QtCore.Signal(list)
  29. shapeMoved = QtCore.Signal()
  30. drawingPolygon = QtCore.Signal(bool)
  31. vertexSelected = QtCore.Signal(bool)
  32. mouseMoved = QtCore.Signal(QtCore.QPointF)
  33. CREATE, EDIT = 0, 1
  34. # polygon, rectangle, line, or point
  35. _createMode = "polygon"
  36. _fill_drawing = False
  37. def __init__(self, *args, **kwargs):
  38. self.epsilon = kwargs.pop("epsilon", 10.0)
  39. self.detection_model_path=kwargs.pop("detection_model_path",None)
  40. self.segmentation_model_path=kwargs.pop("segmentation_model_path",None)
  41. self.decoding_model_path=kwargs.pop("decoding_model_path", None)
  42. self.double_click = kwargs.pop("double_click", "close")
  43. self._characters = load_barcode_dict() # Load dictionary once
  44. logger.debug(f"Loaded barcode dictionary with {len(self._characters)} characters.")
  45. if self.double_click not in [None, "close"]:
  46. raise ValueError(
  47. "Unexpected value for double_click event: {}".format(self.double_click)
  48. )
  49. self.num_backups = kwargs.pop("num_backups", 10)
  50. self._crosshair = kwargs.pop(
  51. "crosshair",
  52. {
  53. "polygon": False,
  54. "rectangle": True,
  55. "circle": False,
  56. "line": False,
  57. "point": False,
  58. "linestrip": False,
  59. "ai_polygon": False,
  60. "ai_mask": False,
  61. },
  62. )
  63. super(Canvas, self).__init__(*args, **kwargs)
  64. # Initialise local state.
  65. self.mode = self.EDIT
  66. self.shapes = []
  67. self.shapesBackups = []
  68. self.current = None
  69. self.selectedShapes = [] # save the selected shapes here
  70. self.selectedShapesCopy = []
  71. # self.line represents:
  72. # - createMode == 'polygon': edge from last point to current
  73. # - createMode == 'rectangle': diagonal line of the rectangle
  74. # - createMode == 'line': the line
  75. # - createMode == 'point': the point
  76. self.line = Shape()
  77. self.prevPoint = QtCore.QPoint()
  78. self.prevMovePoint = QtCore.QPoint()
  79. self.offsets = QtCore.QPoint(), QtCore.QPoint()
  80. self.scale = 1.0
  81. self.pixmap = QtGui.QPixmap()
  82. self.visible = {}
  83. self._hideBackround = False
  84. self.hideBackround = False
  85. self.hShape = None
  86. self.prevhShape = None
  87. self.hVertex = None
  88. self.prevhVertex = None
  89. self.hEdge = None
  90. self.prevhEdge = None
  91. self.movingShape = False
  92. self.snapping = True
  93. self.hShapeIsSelected = False
  94. self._painter = QtGui.QPainter()
  95. self._cursor = CURSOR_DEFAULT
  96. self.draw_pred=False
  97. self.pred_bbox_points=None
  98. self.current_bbox_point=None
  99. self.decoded_barcode = None
  100. # Menus:
  101. # 0: right-click without selection and dragging of shapes
  102. # 1: right-click with selection and dragging of shapes
  103. self.menus = (QtWidgets.QMenu(), QtWidgets.QMenu())
  104. # Set widget options.
  105. self.setMouseTracking(True)
  106. self.setFocusPolicy(QtCore.Qt.WheelFocus)
  107. self._ai_model = None
  108. self._detection_model = None
  109. self._segmentation_model = None
  110. self._decoding_model = None
  111. def fillDrawing(self):
  112. return self._fill_drawing
  113. def setFillDrawing(self, value):
  114. self._fill_drawing = value
  115. @property
  116. def createMode(self):
  117. return self._createMode
  118. @createMode.setter
  119. def createMode(self, value):
  120. if value not in [
  121. "polygon",
  122. "rectangle",
  123. "circle",
  124. "line",
  125. "point",
  126. "linestrip",
  127. "ai_polygon",
  128. "ai_mask",
  129. ]:
  130. raise ValueError("Unsupported createMode: %s" % value)
  131. self._createMode = value
  132. def initializeBarcodeModel(self, detection_model_path, segmentation_model_path=None, decoding_model_path=None):
  133. if not detection_model_path:
  134. raise ValueError("Detection model path is required.")
  135. logger.debug("Initializing only detection model: %r" % "BarcodePredictModel")
  136. self._detection_model = labelme.ai.BarcodeDetectModel(detection_model_path)
  137. if segmentation_model_path:
  138. logger.debug("Initializing barcode detection & Segmentation model: %r" % "BarcodePredictModel")
  139. self._segmentation_model = labelme.ai.BarcodeDetectModel(
  140. detection_model_path, segmentation_model_path
  141. )
  142. if decoding_model_path:
  143. logger.debug("Initializing barcode detection, Segmentation model, Decoding_model: %r" % "BarcodePredictModel")
  144. self._decoding_model = labelme.ai.BarcodeDecodeModel(
  145. decoding_model_path
  146. )
  147. if self.pixmap is None:
  148. logger.warning("Pixmap is not set yet")
  149. return
  150. self._detection_model.set_image(
  151. image=labelme.utils.img_qt_to_arr(self.pixmap.toImage())
  152. )
  153. def initializeAiModel(self, name, weight_path=None):
  154. if self._ai_model is not None:
  155. logger.debug("AI model is already initialized.")
  156. return
  157. if name not in [model.name for model in labelme.ai.MODELS]:
  158. raise ValueError("Unsupported AI model: %s" % name)
  159. model_class = [model for model in labelme.ai.MODELS if model.name == name][0]
  160. logger.debug(f"Initializing AI model: {name}")
  161. self._ai_model = model_class(weight_path)
  162. if self.pixmap is None:
  163. logger.warning("Pixmap is not set yet")
  164. return
  165. self._ai_model.set_image(
  166. image=labelme.utils.img_qt_to_arr(self.pixmap.toImage())
  167. # image=self.pixmap.toImage()
  168. )
  169. def storeShapes(self):
  170. shapesBackup = []
  171. for shape in self.shapes:
  172. shapesBackup.append(shape.copy())
  173. if len(self.shapesBackups) > self.num_backups:
  174. self.shapesBackups = self.shapesBackups[-self.num_backups - 1 :]
  175. self.shapesBackups.append(shapesBackup)
  176. @property
  177. def isShapeRestorable(self):
  178. # We save the state AFTER each edit (not before) so for an
  179. # edit to be undoable, we expect the CURRENT and the PREVIOUS state
  180. # to be in the undo stack.
  181. if len(self.shapesBackups) < 2:
  182. return False
  183. return True
  184. def restoreShape(self):
  185. # This does _part_ of the job of restoring shapes.
  186. # The complete process is also done in app.py::undoShapeEdit
  187. # and app.py::loadShapes and our own Canvas::loadShapes function.
  188. if not self.isShapeRestorable:
  189. return
  190. # print(f"shape is restorable")
  191. self.shapesBackups.pop() # latest
  192. # The application will eventually call Canvas.loadShapes which will
  193. # push this right back onto the stack.
  194. shapesBackup = self.shapesBackups.pop()
  195. self.shapes = shapesBackup
  196. self.selectedShapes = []
  197. for shape in self.shapes:
  198. shape.selected = False
  199. self.update()
  200. def enterEvent(self, ev):
  201. self.overrideCursor(self._cursor)
  202. def leaveEvent(self, ev):
  203. self.unHighlight()
  204. self.restoreCursor()
  205. def focusOutEvent(self, ev):
  206. self.restoreCursor()
  207. def isVisible(self, shape):
  208. return self.visible.get(shape, True)
  209. def drawing(self):
  210. return self.mode == self.CREATE
  211. def editing(self):
  212. return self.mode == self.EDIT
  213. def setEditing(self, value=True):
  214. self.mode = self.EDIT if value else self.CREATE
  215. if self.mode == self.EDIT:
  216. # CREATE -> EDIT
  217. self.repaint() # clear crosshair
  218. else:
  219. # EDIT -> CREATE
  220. self.unHighlight()
  221. self.deSelectShape()
  222. def unHighlight(self):
  223. if self.hShape:
  224. self.hShape.highlightClear()
  225. self.update()
  226. self.prevhShape = self.hShape
  227. self.prevhVertex = self.hVertex
  228. self.prevhEdge = self.hEdge
  229. self.hShape = self.hVertex = self.hEdge = None
  230. def selectedVertex(self):
  231. return self.hVertex is not None
  232. def selectedEdge(self):
  233. return self.hEdge is not None
  234. def mouseMoveEvent(self, ev):
  235. """Update line with last point and current coordinates."""
  236. try:
  237. if QT5:
  238. pos = self.transformPos(ev.localPos())
  239. else:
  240. pos = self.transformPos(ev.posF())
  241. except AttributeError:
  242. return
  243. self.mouseMoved.emit(pos)
  244. self.prevMovePoint = pos
  245. self.restoreCursor()
  246. is_shift_pressed = ev.modifiers() & QtCore.Qt.ShiftModifier
  247. # Polygon drawing.
  248. if self.drawing():
  249. if self.createMode in ["ai_polygon", "ai_mask"]:
  250. self.line.shape_type = "points"
  251. else:
  252. self.line.shape_type = self.createMode
  253. self.overrideCursor(CURSOR_DRAW)
  254. if not self.current:
  255. self.repaint() # draw crosshair
  256. return
  257. if self.outOfPixmap(pos):
  258. # Don't allow the user to draw outside the pixmap.
  259. # Project the point to the pixmap's edges.
  260. pos = self.intersectionPoint(self.current[-1], pos)
  261. elif (
  262. self.snapping
  263. and len(self.current) > 1
  264. and self.createMode == "polygon"
  265. and self.closeEnough(pos, self.current[0])
  266. ):
  267. # Attract line to starting point and
  268. # colorise to alert the user.
  269. pos = self.current[0]
  270. self.overrideCursor(CURSOR_POINT)
  271. self.current.highlightVertex(0, Shape.NEAR_VERTEX)
  272. if self.createMode in ["polygon", "linestrip"]:
  273. self.line.points = [self.current[-1], pos]
  274. self.line.point_labels = [1, 1]
  275. elif self.createMode in ["ai_polygon", "ai_mask"]:
  276. self.line.points = [self.current.points[-1], pos]
  277. self.line.point_labels = [
  278. self.current.point_labels[-1],
  279. 0 if is_shift_pressed else 1,
  280. ]
  281. elif self.createMode == "rectangle":
  282. self.line.points = [self.current[0], pos]
  283. self.line.point_labels = [1, 1]
  284. self.line.close()
  285. elif self.createMode == "circle":
  286. self.line.points = [self.current[0], pos]
  287. self.line.point_labels = [1, 1]
  288. self.line.shape_type = "circle"
  289. elif self.createMode == "line":
  290. self.line.points = [self.current[0], pos]
  291. self.line.point_labels = [1, 1]
  292. self.line.close()
  293. elif self.createMode == "point":
  294. self.line.points = [self.current[0]]
  295. self.line.point_labels = [1]
  296. self.line.close()
  297. # print("self.line.points", self.line.points)
  298. assert len(self.line.points) == len(self.line.point_labels)
  299. self.repaint()
  300. self.current.highlightClear()
  301. return
  302. # Polygon copy moving.
  303. if QtCore.Qt.RightButton & ev.buttons():
  304. if self.selectedShapesCopy and self.prevPoint:
  305. self.overrideCursor(CURSOR_MOVE)
  306. self.boundedMoveShapes(self.selectedShapesCopy, pos)
  307. self.repaint()
  308. elif self.selectedShapes:
  309. self.selectedShapesCopy = [s.copy() for s in self.selectedShapes]
  310. self.repaint()
  311. return
  312. # Polygon/Vertex moving.
  313. if QtCore.Qt.LeftButton & ev.buttons():
  314. if self.selectedVertex():
  315. self.boundedMoveVertex(pos)
  316. self.repaint()
  317. self.movingShape = True
  318. elif self.selectedShapes and self.prevPoint:
  319. self.overrideCursor(CURSOR_MOVE)
  320. self.boundedMoveShapes(self.selectedShapes, pos)
  321. self.repaint()
  322. self.movingShape = True
  323. return
  324. # Just hovering over the canvas, 2 possibilities:
  325. # - Highlight shapes
  326. # - Highlight vertex
  327. # Update shape/vertex fill and tooltip value accordingly.
  328. self.setToolTip(self.tr("Image"))
  329. for shape in reversed([s for s in self.shapes if self.isVisible(s)]):
  330. # Look for a nearby vertex to highlight. If that fails,
  331. # check if we happen to be inside a shape.
  332. index = shape.nearestVertex(pos, self.epsilon / self.scale)
  333. index_edge = shape.nearestEdge(pos, self.epsilon / self.scale)
  334. if index is not None:
  335. if self.selectedVertex():
  336. self.hShape.highlightClear()
  337. self.prevhVertex = self.hVertex = index
  338. self.prevhShape = self.hShape = shape
  339. self.prevhEdge = self.hEdge
  340. self.hEdge = None
  341. shape.highlightVertex(index, shape.MOVE_VERTEX)
  342. self.overrideCursor(CURSOR_POINT)
  343. self.setToolTip(self.tr("Click & drag to move point"))
  344. self.setStatusTip(self.toolTip())
  345. self.update()
  346. break
  347. elif index_edge is not None and shape.canAddPoint():
  348. if self.selectedVertex():
  349. self.hShape.highlightClear()
  350. self.prevhVertex = self.hVertex
  351. self.hVertex = None
  352. self.prevhShape = self.hShape = shape
  353. self.prevhEdge = self.hEdge = index_edge
  354. self.overrideCursor(CURSOR_POINT)
  355. self.setToolTip(self.tr("Click to create point"))
  356. self.setStatusTip(self.toolTip())
  357. self.update()
  358. break
  359. elif len(shape.points)!=0 and shape.containsPoint(pos) :
  360. if self.selectedVertex():
  361. self.hShape.highlightClear()
  362. self.prevhVertex = self.hVertex
  363. self.hVertex = None
  364. self.prevhShape = self.hShape = shape
  365. self.prevhEdge = self.hEdge
  366. self.hEdge = None
  367. self.setToolTip(
  368. self.tr("Click & drag to move shape '%s'") % shape.label
  369. )
  370. self.setStatusTip(self.toolTip())
  371. self.overrideCursor(CURSOR_GRAB)
  372. self.update()
  373. break
  374. else: # Nothing found, clear highlights, reset state.
  375. self.unHighlight()
  376. self.vertexSelected.emit(self.hVertex is not None)
  377. def addPointToEdge(self):
  378. shape = self.prevhShape
  379. index = self.prevhEdge
  380. point = self.prevMovePoint
  381. if shape is None or index is None or point is None:
  382. return
  383. shape.insertPoint(index, point)
  384. shape.highlightVertex(index, shape.MOVE_VERTEX)
  385. self.hShape = shape
  386. self.hVertex = index
  387. self.hEdge = None
  388. self.movingShape = True
  389. def removeSelectedPoint(self):
  390. shape = self.prevhShape
  391. index = self.prevhVertex
  392. if shape is None or index is None:
  393. return
  394. shape.removePoint(index)
  395. shape.highlightClear()
  396. self.hShape = shape
  397. self.prevhVertex = None
  398. self.movingShape = True # Save changes
  399. def mousePressEvent(self, ev):
  400. if QT5:
  401. pos = self.transformPos(ev.localPos())
  402. else:
  403. pos = self.transformPos(ev.posF())
  404. is_shift_pressed = ev.modifiers() & QtCore.Qt.ShiftModifier
  405. if ev.button() == QtCore.Qt.LeftButton:
  406. if self.drawing():
  407. if self.current:
  408. # Add point to existing shape.
  409. if self.createMode == "polygon":
  410. self.current.addPoint(self.line[1])
  411. self.line[0] = self.current[-1]
  412. if self.current.isClosed():
  413. self.finalise()
  414. elif self.createMode in ["rectangle", "circle", "line"]:
  415. assert len(self.current.points) == 1
  416. self.current.points = self.line.points
  417. self.finalise()
  418. elif self.createMode == "linestrip":
  419. self.current.addPoint(self.line[1])
  420. self.line[0] = self.current[-1]
  421. if int(ev.modifiers()) == QtCore.Qt.ControlModifier:
  422. self.finalise()
  423. elif self.createMode in ["ai_polygon", "ai_mask"]:
  424. self.current.addPoint(
  425. self.line.points[1],
  426. label=self.line.point_labels[1],
  427. )
  428. self.line.points[0] = self.current.points[-1]
  429. self.line.point_labels[0] = self.current.point_labels[-1]
  430. if ev.modifiers() & QtCore.Qt.ControlModifier:
  431. self.finalise()
  432. elif not self.outOfPixmap(pos):
  433. # Create new shape.
  434. self.current = Shape(
  435. shape_type="points"
  436. if self.createMode in ["ai_polygon", "ai_mask"]
  437. else self.createMode
  438. )
  439. self.current.addPoint(pos, label=0 if is_shift_pressed else 1)
  440. if self.createMode == "point":
  441. self.finalise()
  442. elif (
  443. self.createMode in ["ai_polygon", "ai_mask"]
  444. and ev.modifiers() & QtCore.Qt.ControlModifier
  445. ):
  446. self.finalise()
  447. else:
  448. if self.createMode == "circle":
  449. self.current.shape_type = "circle"
  450. self.line.points = [pos, pos]
  451. if (
  452. self.createMode in ["ai_polygon", "ai_mask"]
  453. and is_shift_pressed
  454. ):
  455. self.line.point_labels = [0, 0]
  456. else:
  457. self.line.point_labels = [1, 1]
  458. self.setHiding()
  459. self.drawingPolygon.emit(True)
  460. self.update()
  461. elif self.editing():
  462. if self.selectedEdge():
  463. self.addPointToEdge()
  464. elif (
  465. self.selectedVertex()
  466. and int(ev.modifiers()) == QtCore.Qt.ShiftModifier
  467. ):
  468. # Delete point if: left-click + SHIFT on a point
  469. self.removeSelectedPoint()
  470. group_mode = int(ev.modifiers()) == QtCore.Qt.ControlModifier
  471. self.selectShapePoint(pos, multiple_selection_mode=group_mode)
  472. self.prevPoint = pos
  473. self.repaint()
  474. elif ev.button() == QtCore.Qt.RightButton and self.editing():
  475. group_mode = int(ev.modifiers()) == QtCore.Qt.ControlModifier
  476. if not self.selectedShapes or (
  477. self.hShape is not None and self.hShape not in self.selectedShapes
  478. ):
  479. self.selectShapePoint(pos, multiple_selection_mode=group_mode)
  480. self.repaint()
  481. self.prevPoint = pos
  482. def mouseReleaseEvent(self, ev):
  483. if ev.button() == QtCore.Qt.RightButton:
  484. menu = self.menus[len(self.selectedShapesCopy) > 0]
  485. self.restoreCursor()
  486. if not menu.exec_(self.mapToGlobal(ev.pos())) and self.selectedShapesCopy:
  487. # Cancel the move by deleting the shadow copy.
  488. self.selectedShapesCopy = []
  489. self.repaint()
  490. elif ev.button() == QtCore.Qt.LeftButton:
  491. if self.editing():
  492. if (
  493. self.hShape is not None
  494. and self.hShapeIsSelected
  495. and not self.movingShape
  496. ):
  497. self.selectionChanged.emit(
  498. [x for x in self.selectedShapes if x != self.hShape]
  499. )
  500. if self.movingShape and self.hShape:
  501. index = self.shapes.index(self.hShape)
  502. if self.shapesBackups[-1][index].points != self.shapes[index].points:
  503. self.storeShapes()
  504. self.shapeMoved.emit()
  505. self.movingShape = False
  506. def endMove(self, copy):
  507. assert self.selectedShapes and self.selectedShapesCopy
  508. assert len(self.selectedShapesCopy) == len(self.selectedShapes)
  509. if copy:
  510. for i, shape in enumerate(self.selectedShapesCopy):
  511. self.shapes.append(shape)
  512. self.selectedShapes[i].selected = False
  513. self.selectedShapes[i] = shape
  514. else:
  515. for i, shape in enumerate(self.selectedShapesCopy):
  516. self.selectedShapes[i].points = shape.points
  517. self.selectedShapesCopy = []
  518. self.repaint()
  519. self.storeShapes()
  520. return True
  521. def hideBackroundShapes(self, value):
  522. self.hideBackround = value
  523. if self.selectedShapes:
  524. # Only hide other shapes if there is a current selection.
  525. # Otherwise the user will not be able to select a shape.
  526. self.setHiding(True)
  527. self.update()
  528. def setHiding(self, enable=True):
  529. self._hideBackround = self.hideBackround if enable else False
  530. def canCloseShape(self):
  531. return self.drawing() and (
  532. (self.current and len(self.current) > 2)
  533. or self.createMode in ["ai_polygon", "ai_mask"]
  534. )
  535. def mouseDoubleClickEvent(self, ev):
  536. if self.double_click != "close":
  537. return
  538. if (
  539. self.createMode == "polygon" and self.canCloseShape()
  540. ) or self.createMode in ["ai_polygon", "ai_mask"]:
  541. self.finalise()
  542. def selectShapes(self, shapes):
  543. self.setHiding()
  544. self.selectionChanged.emit(shapes)
  545. self.update()
  546. def selectShapePoint(self, point, multiple_selection_mode):
  547. """Select the first shape created which contains this point."""
  548. if self.selectedVertex(): # A vertex is marked for selection.
  549. index, shape = self.hVertex, self.hShape
  550. shape.highlightVertex(index, shape.MOVE_VERTEX)
  551. else:
  552. for shape in reversed(self.shapes):
  553. if self.isVisible(shape) and shape.containsPoint(point):
  554. self.setHiding()
  555. if shape not in self.selectedShapes:
  556. if multiple_selection_mode:
  557. self.selectionChanged.emit(self.selectedShapes + [shape])
  558. else:
  559. self.selectionChanged.emit([shape])
  560. self.hShapeIsSelected = False
  561. else:
  562. self.hShapeIsSelected = True
  563. self.calculateOffsets(point)
  564. return
  565. self.deSelectShape()
  566. def calculateOffsets(self, point):
  567. left = self.pixmap.width() - 1
  568. right = 0
  569. top = self.pixmap.height() - 1
  570. bottom = 0
  571. for s in self.selectedShapes:
  572. rect = s.boundingRect()
  573. if rect.left() < left:
  574. left = rect.left()
  575. if rect.right() > right:
  576. right = rect.right()
  577. if rect.top() < top:
  578. top = rect.top()
  579. if rect.bottom() > bottom:
  580. bottom = rect.bottom()
  581. x1 = left - point.x()
  582. y1 = top - point.y()
  583. x2 = right - point.x()
  584. y2 = bottom - point.y()
  585. self.offsets = QtCore.QPointF(x1, y1), QtCore.QPointF(x2, y2)
  586. def boundedMoveVertex(self, pos):
  587. index, shape = self.hVertex, self.hShape
  588. point = shape[index]
  589. if self.outOfPixmap(pos):
  590. pos = self.intersectionPoint(point, pos)
  591. shape.moveVertexBy(index, pos - point)
  592. def boundedMoveShapes(self, shapes, pos):
  593. if self.outOfPixmap(pos):
  594. return False # No need to move
  595. o1 = pos + self.offsets[0]
  596. if self.outOfPixmap(o1):
  597. pos -= QtCore.QPointF(min(0, o1.x()), min(0, o1.y()))
  598. o2 = pos + self.offsets[1]
  599. if self.outOfPixmap(o2):
  600. pos += QtCore.QPointF(
  601. min(0, self.pixmap.width() - o2.x()),
  602. min(0, self.pixmap.height() - o2.y()),
  603. )
  604. # XXX: The next line tracks the new position of the cursor
  605. # relative to the shape, but also results in making it
  606. # a bit "shaky" when nearing the border and allows it to
  607. # go outside of the shape's area for some reason.
  608. # self.calculateOffsets(self.selectedShapes, pos)
  609. dp = pos - self.prevPoint
  610. if dp:
  611. for shape in shapes:
  612. shape.moveBy(dp)
  613. self.prevPoint = pos
  614. return True
  615. return False
  616. def deSelectShape(self):
  617. if self.selectedShapes:
  618. self.setHiding(False)
  619. self.selectionChanged.emit([])
  620. self.hShapeIsSelected = False
  621. self.update()
  622. def deleteSelected(self):
  623. deleted_shapes = []
  624. if self.selectedShapes:
  625. for shape in self.selectedShapes:
  626. self.shapes.remove(shape)
  627. deleted_shapes.append(shape)
  628. self.storeShapes()
  629. self.selectedShapes = []
  630. self.update()
  631. return deleted_shapes
  632. def deleteShape(self, shape):
  633. if shape in self.selectedShapes:
  634. self.selectedShapes.remove(shape)
  635. if shape in self.shapes:
  636. self.shapes.remove(shape)
  637. self.storeShapes()
  638. self.update()
  639. def duplicateSelectedShapes(self):
  640. if self.selectedShapes:
  641. self.selectedShapesCopy = [s.copy() for s in self.selectedShapes]
  642. self.boundedShiftShapes(self.selectedShapesCopy)
  643. self.endMove(copy=True)
  644. return self.selectedShapes
  645. def boundedShiftShapes(self, shapes):
  646. # Try to move in one direction, and if it fails in another.
  647. # Give up if both fail.
  648. point = shapes[0][0]
  649. offset = QtCore.QPointF(2.0, 2.0)
  650. self.offsets = QtCore.QPoint(), QtCore.QPoint()
  651. self.prevPoint = point
  652. if not self.boundedMoveShapes(shapes, point - offset):
  653. self.boundedMoveShapes(shapes, point + offset)
  654. def paintEvent(self, event):
  655. if not self.pixmap:
  656. return super(Canvas, self).paintEvent(event)
  657. p = self._painter
  658. p.begin(self)
  659. p.setRenderHint(QtGui.QPainter.Antialiasing)
  660. p.setRenderHint(QtGui.QPainter.HighQualityAntialiasing)
  661. p.setRenderHint(QtGui.QPainter.SmoothPixmapTransform)
  662. p.scale(self.scale, self.scale)
  663. p.translate(self.offsetToCenter())
  664. p.drawPixmap(0, 0, self.pixmap)
  665. # draw crosshair
  666. if (
  667. self._crosshair[self._createMode]
  668. and self.drawing()
  669. and self.prevMovePoint
  670. and not self.outOfPixmap(self.prevMovePoint)
  671. ):
  672. p.setPen(QtGui.QColor(0, 0, 0))
  673. p.drawLine(
  674. 0,
  675. int(self.prevMovePoint.y()),
  676. self.width() - 1,
  677. int(self.prevMovePoint.y()),
  678. )
  679. p.drawLine(
  680. int(self.prevMovePoint.x()),
  681. 0,
  682. int(self.prevMovePoint.x()),
  683. self.height() - 1,
  684. )
  685. Shape.scale = self.scale
  686. for shape in self.shapes:
  687. if (shape.selected or not self._hideBackround) and self.isVisible(shape):
  688. shape.fill = shape.selected or shape == self.hShape
  689. shape.paint(p)
  690. if self.current:
  691. self.current.paint(p)
  692. assert len(self.line.points) == len(self.line.point_labels)
  693. self.line.paint(p)
  694. if self.selectedShapesCopy:
  695. for s in self.selectedShapesCopy:
  696. s.paint(p)
  697. # if(self.draw_pred and self.current is not None):
  698. # print("pred mode on")
  699. # for bbox_points in self.pred_bbox_points:
  700. # drawing_shape = self.current.copy()
  701. # drawing_shape.setShapeRefined(
  702. # shape_type="polygon",
  703. # points=[QtCore.QPointF(point[0], point[1]) for point in bbox_points],
  704. # point_labels=[1]*len(bbox_points)
  705. # )
  706. # drawing_shape.fill = self.fillDrawing()
  707. # drawing_shape.selected = True
  708. # drawing_shape.paint(p)
  709. # self.draw_pred=False
  710. if (
  711. self.fillDrawing()
  712. and self.createMode == "polygon"
  713. and self.current is not None
  714. and len(self.current.points) >= 2
  715. ):
  716. drawing_shape = self.current.copy()
  717. if drawing_shape.fill_color.getRgb()[3] == 0:
  718. logger.warning(
  719. "fill_drawing=true, but fill_color is transparent,"
  720. " so forcing to be opaque."
  721. )
  722. drawing_shape.fill_color.setAlpha(64)
  723. drawing_shape.addPoint(self.line[1])
  724. drawing_shape.fill = True
  725. drawing_shape.paint(p)
  726. elif self.createMode == "ai_polygon" and self.current is not None:
  727. drawing_shape = self.current.copy()
  728. drawing_shape.addPoint(
  729. point=self.line.points[1],
  730. label=self.line.point_labels[1],
  731. )
  732. points = self._detection_model.predict_polygon_from_points(
  733. points=[[point.x(), point.y()] for point in drawing_shape.points],
  734. point_labels=drawing_shape.point_labels,
  735. )
  736. if len(points) > 2:
  737. drawing_shape.setShapeRefined(
  738. shape_type="polygon",
  739. points=[QtCore.QPointF(point[0], point[1]) for point in points],
  740. point_labels=[1] * len(points),
  741. )
  742. drawing_shape.fill = self.fillDrawing()
  743. drawing_shape.selected = True
  744. drawing_shape.paint(p)
  745. elif self.createMode == "ai_mask" and self.current is not None:
  746. drawing_shape = self.current.copy()
  747. drawing_shape.addPoint(
  748. point=self.line.points[1],
  749. label=self.line.point_labels[1],
  750. )
  751. mask = self._detection_model.predict_mask_from_points(
  752. points=[[point.x(), point.y()] for point in drawing_shape.points],
  753. point_labels=drawing_shape.point_labels,
  754. )
  755. y1, x1, y2, x2 = imgviz.instances.masks_to_bboxes([mask])[0].astype(int)
  756. drawing_shape.setShapeRefined(
  757. shape_type="mask",
  758. points=[QtCore.QPointF(x1, y1), QtCore.QPointF(x2, y2)],
  759. point_labels=[1, 1],
  760. mask=mask[y1 : y2 + 1, x1 : x2 + 1],
  761. )
  762. drawing_shape.selected = True
  763. drawing_shape.paint(p)
  764. p.end()
  765. def transformPos(self, point):
  766. """Convert from widget-logical coordinates to painter-logical ones."""
  767. return point / self.scale - self.offsetToCenter()
  768. def offsetToCenter(self):
  769. s = self.scale
  770. area = super(Canvas, self).size()
  771. w, h = self.pixmap.width() * s, self.pixmap.height() * s
  772. aw, ah = area.width(), area.height()
  773. x = (aw - w) / (2 * s) if aw > w else 0
  774. y = (ah - h) / (2 * s) if ah > h else 0
  775. return QtCore.QPointF(x, y)
  776. def outOfPixmap(self, p):
  777. w, h = self.pixmap.width(), self.pixmap.height()
  778. return not (0 <= p.x() <= w - 1 and 0 <= p.y() <= h - 1)
  779. def finalise(self):
  780. if(self.current is None):
  781. return
  782. # If in manual mode and points are finalized
  783. if self.createMode in ["polygon", "rectangle"] and len(self.current.points) > 0:
  784. if self._detection_model is None:
  785. logger.info(f"Initializing AI model")
  786. self.initializeBarcodeModel(self.detection_model_path, self.segmentation_model_path, self.decoding_model_path)
  787. # Extract the points from the manually drawn shape
  788. manual_points = [[point.x(), point.y()] for point in self.current.points]
  789. print(manual_points)
  790. logger.debug(f"Manually drawn points: {manual_points}")
  791. self.detect_and_segment(manual_points=manual_points)
  792. self.mode = self.EDIT # Or any appropriate mode to disable drawing behavior
  793. return
  794. if self.createMode == "ai_polygon":
  795. # convert points to polygon by an AI model
  796. assert self.current.shape_type == "points"
  797. points = self._detection_model.predict_polygon_from_points(
  798. points=[[point.x(), point.y()] for point in self.current.points],
  799. point_labels=self.current.point_labels,
  800. )
  801. self.current.setShapeRefined(
  802. points=[QtCore.QPointF(point[0], point[1]) for point in points],
  803. point_labels=[1] * len(points),
  804. shape_type="polygon",
  805. )
  806. elif self.createMode == "ai_mask":
  807. # convert points to mask by an AI model
  808. assert self.current.shape_type == "points"
  809. mask = self._detection_model.predict_mask_from_points(
  810. points=[[point.x(), point.y()] for point in self.current.points],
  811. point_labels=self.current.point_labels,
  812. )
  813. y1, x1, y2, x2 = imgviz.instances.masks_to_bboxes([mask])[0].astype(int)
  814. self.current.setShapeRefined(
  815. shape_type="mask",
  816. points=[QtCore.QPointF(x1, y1), QtCore.QPointF(x2, y2)],
  817. point_labels=[1, 1],
  818. mask=mask[y1 : y2 + 1, x1 : x2 + 1],
  819. )
  820. elif self.pred_bbox_points is not None and self.draw_pred:
  821. print("pred mode on")
  822. current_copy=self.current.copy()
  823. for bbox_point in self.pred_bbox_points:
  824. drawing_shape=current_copy.copy()
  825. drawing_shape.setShapeRefined(
  826. shape_type="polygon",
  827. points=[QtCore.QPointF(point[0], point[1]) for point in bbox_point],
  828. point_labels=[1]*len(bbox_point)
  829. )
  830. drawing_shape.description = self.decoded_barcode
  831. drawing_shape.close()
  832. self.shapes.append(drawing_shape)
  833. self.storeShapes()
  834. self.update()
  835. self.newShape.emit()
  836. # self.pred_bbox_points = None
  837. # self.draw_pred = False
  838. # self.detect_and_segment() # Process the next detection
  839. current_copy.close()
  840. current_copy=None
  841. if(self.current):
  842. self.current.close()
  843. self.current = None
  844. self.setHiding(False)
  845. self.draw_pred=False
  846. return
  847. self.current.close()
  848. self.shapes.append(self.current)
  849. self.storeShapes()
  850. self.current = None
  851. self.setHiding(False)
  852. self.newShape.emit()
  853. self.update()
  854. def closeEnough(self, p1, p2):
  855. # d = distance(p1 - p2)
  856. # m = (p1-p2).manhattanLength()
  857. # print "d %.2f, m %d, %.2f" % (d, m, d - m)
  858. # divide by scale to allow more precision when zoomed in
  859. return labelme.utils.distance(p1 - p2) < (self.epsilon / self.scale)
  860. def intersectionPoint(self, p1, p2):
  861. # Cycle through each image edge in clockwise fashion,
  862. # and find the one intersecting the current line segment.
  863. # http://paulbourke.net/geometry/lineline2d/
  864. size = self.pixmap.size()
  865. points = [
  866. (0, 0),
  867. (size.width() - 1, 0),
  868. (size.width() - 1, size.height() - 1),
  869. (0, size.height() - 1),
  870. ]
  871. # x1, y1 should be in the pixmap, x2, y2 should be out of the pixmap
  872. x1 = min(max(p1.x(), 0), size.width() - 1)
  873. y1 = min(max(p1.y(), 0), size.height() - 1)
  874. x2, y2 = p2.x(), p2.y()
  875. d, i, (x, y) = min(self.intersectingEdges((x1, y1), (x2, y2), points))
  876. x3, y3 = points[i]
  877. x4, y4 = points[(i + 1) % 4]
  878. if (x, y) == (x1, y1):
  879. # Handle cases where previous point is on one of the edges.
  880. if x3 == x4:
  881. return QtCore.QPointF(x3, min(max(0, y2), max(y3, y4)))
  882. else: # y3 == y4
  883. return QtCore.QPointF(min(max(0, x2), max(x3, x4)), y3)
  884. return QtCore.QPointF(x, y)
  885. def intersectingEdges(self, point1, point2, points):
  886. """Find intersecting edges.
  887. For each edge formed by `points', yield the intersection
  888. with the line segment `(x1,y1) - (x2,y2)`, if it exists.
  889. Also return the distance of `(x2,y2)' to the middle of the
  890. edge along with its index, so that the one closest can be chosen.
  891. """
  892. (x1, y1) = point1
  893. (x2, y2) = point2
  894. for i in range(4):
  895. x3, y3 = points[i]
  896. x4, y4 = points[(i + 1) % 4]
  897. denom = (y4 - y3) * (x2 - x1) - (x4 - x3) * (y2 - y1)
  898. nua = (x4 - x3) * (y1 - y3) - (y4 - y3) * (x1 - x3)
  899. nub = (x2 - x1) * (y1 - y3) - (y2 - y1) * (x1 - x3)
  900. if denom == 0:
  901. # This covers two cases:
  902. # nua == nub == 0: Coincident
  903. # otherwise: Parallel
  904. continue
  905. ua, ub = nua / denom, nub / denom
  906. if 0 <= ua <= 1 and 0 <= ub <= 1:
  907. x = x1 + ua * (x2 - x1)
  908. y = y1 + ua * (y2 - y1)
  909. m = QtCore.QPointF((x3 + x4) / 2, (y3 + y4) / 2)
  910. d = labelme.utils.distance(m - QtCore.QPointF(x2, y2))
  911. yield d, i, (x, y)
  912. # These two, along with a call to adjustSize are required for the
  913. # scroll area.
  914. def sizeHint(self):
  915. return self.minimumSizeHint()
  916. def minimumSizeHint(self):
  917. if self.pixmap:
  918. return self.scale * self.pixmap.size()
  919. return super(Canvas, self).minimumSizeHint()
  920. def wheelEvent(self, ev):
  921. if QT5:
  922. mods = ev.modifiers()
  923. delta = ev.angleDelta()
  924. if QtCore.Qt.ControlModifier == int(mods):
  925. # with Ctrl/Command key
  926. # zoom
  927. self.zoomRequest.emit(delta.y(), ev.pos())
  928. else:
  929. # scroll
  930. self.scrollRequest.emit(delta.x(), QtCore.Qt.Horizontal)
  931. self.scrollRequest.emit(delta.y(), QtCore.Qt.Vertical)
  932. else:
  933. if ev.orientation() == QtCore.Qt.Vertical:
  934. mods = ev.modifiers()
  935. if QtCore.Qt.ControlModifier == int(mods):
  936. # with Ctrl/Command key
  937. self.zoomRequest.emit(ev.delta(), ev.pos())
  938. else:
  939. self.scrollRequest.emit(
  940. ev.delta(),
  941. QtCore.Qt.Horizontal
  942. if (QtCore.Qt.ShiftModifier == int(mods))
  943. else QtCore.Qt.Vertical,
  944. )
  945. else:
  946. self.scrollRequest.emit(ev.delta(), QtCore.Qt.Horizontal)
  947. ev.accept()
  948. def moveByKeyboard(self, offset):
  949. if self.selectedShapes:
  950. self.boundedMoveShapes(self.selectedShapes, self.prevPoint + offset)
  951. self.repaint()
  952. self.movingShape = True
  953. def keyPressEvent(self, ev):
  954. modifiers = ev.modifiers()
  955. key = ev.key()
  956. if self.drawing():
  957. if key == QtCore.Qt.Key_Escape and self.current:
  958. self.current = None
  959. self.drawingPolygon.emit(False)
  960. self.update()
  961. elif key == QtCore.Qt.Key_Return and self.canCloseShape():
  962. self.finalise()
  963. elif modifiers == QtCore.Qt.AltModifier:
  964. self.snapping = False
  965. elif key == QtCore.Qt.Key_V:
  966. if self._detection_model is None:
  967. logger.info(f"Initializing AI model")
  968. self.initializeBarcodeModel(self.detection_model_path, self.segmentation_model_path)
  969. self.current = Shape(
  970. shape_type="points" if self.createMode in ["ai_polygon", "ai_mask"] else self.createMode
  971. )
  972. if self._detection_model:
  973. if self.segmentation_model_path is None:
  974. logger.info(f"Performing detection only.")
  975. # Get prediction from model
  976. self.pred_bbox_points = self._detection_model.predict_polygon_from_points()
  977. print("Predicted Bounding Box Points:", self.pred_bbox_points)
  978. if self.pred_bbox_points:
  979. self.draw_pred = True
  980. self.finalise()
  981. else:
  982. print("No bounding boxes detected.")
  983. else:
  984. # logger.info(f"Performing detection, segmentation and decoding.")
  985. self.initializeBarcodeModel(self.detection_model_path, self.segmentation_model_path, self.decoding_model_path)
  986. self.detect_and_segment(manual_points=None)
  987. elif self.editing():
  988. if key == QtCore.Qt.Key_Up:
  989. self.moveByKeyboard(QtCore.QPointF(0.0, -MOVE_SPEED))
  990. elif key == QtCore.Qt.Key_Down:
  991. self.moveByKeyboard(QtCore.QPointF(0.0, MOVE_SPEED))
  992. elif key == QtCore.Qt.Key_Left:
  993. self.moveByKeyboard(QtCore.QPointF(-MOVE_SPEED, 0.0))
  994. elif key == QtCore.Qt.Key_Right:
  995. self.moveByKeyboard(QtCore.QPointF(MOVE_SPEED, 0.0))
  996. def scale_points(self, points, mask_shape, cropped_shape, x_min, y_min):
  997. """
  998. Scale the points from the segmentation mask space back to the original cropped image coordinates.
  999. Args:
  1000. points (np.ndarray): Points to scale (Nx2 array or list of lists).
  1001. mask_shape (tuple): Shape of the segmentation mask (height, width).
  1002. cropped_shape (tuple): Shape of the cropped image in the original image (height, width).
  1003. x_min (int): Minimum x-coordinate of the cropped region in the original image.
  1004. y_min (int): Minimum y-coordinate of the cropped region in the original image.
  1005. Returns:
  1006. scaled_points (list): List of scaled points as [[x1, y1], [x2, y2], ...].
  1007. """
  1008. # Compute scaling factors
  1009. scale_x = cropped_shape[1] / mask_shape[1] # Scale factor for x-axis
  1010. scale_y = cropped_shape[0] / mask_shape[0] # Scale factor for y-axis
  1011. # Scale and translate the points
  1012. scaled_points = []
  1013. for pt in points:
  1014. scaled_x = int(pt[0] * scale_x) + x_min
  1015. scaled_y = int(pt[1] * scale_y) + y_min
  1016. scaled_points.append([scaled_x, scaled_y])
  1017. # logger.debug(f"scaled points: {scaled_points}")
  1018. return scaled_points
  1019. def expand_bbox(self, x_min, y_min, x_max, y_max, factor=1.3):
  1020. """
  1021. Expands the bounding box by a given factor.
  1022. """
  1023. center_x = (x_max + x_min) / 2
  1024. center_y = (y_max + y_min) / 2
  1025. width = (x_max - x_min) * factor
  1026. height = (y_max - y_min) * factor
  1027. new_x_min = int(center_x - width / 2)
  1028. new_y_min = int(center_y - height / 2)
  1029. new_x_max = int(center_x + width / 2)
  1030. new_y_max = int(center_y + height / 2)
  1031. # Ensure the coordinates do not go out of image bounds
  1032. new_x_min = max(0, new_x_min)
  1033. new_y_min = max(0, new_y_min)
  1034. new_x_max = min(self.pixmap.width() - 1, new_x_max)
  1035. new_y_max = min(self.pixmap.height() - 1, new_y_max)
  1036. return new_x_min, new_y_min, new_x_max, new_y_max
  1037. def detect_and_segment(self, manual_points=None):
  1038. """
  1039. Perform detection and segmentation (if both models are available).
  1040. """
  1041. logger.info("Performing detection and segmentation.")
  1042. self.current = Shape(
  1043. shape_type="points" if self.createMode in ["ai_polygon", "ai_mask"] else self.createMode
  1044. )
  1045. print("self.current",self.current)
  1046. # Step 1: detection bounding box points
  1047. detection_results = [manual_points] if manual_points else self._detection_model.predict_polygon_from_points()
  1048. print(detection_results)
  1049. if not detection_results or len(detection_results) == 0:
  1050. logger.warning("No detection found")
  1051. return
  1052. logger.debug(f"Detection results: {detection_results}")
  1053. all_segmentation_results = []
  1054. rotated = False
  1055. # Step 2: Loop through each detection result since there are multiple per image
  1056. for detection_idx, detection_result in enumerate(detection_results):
  1057. logger.debug(f"Processing detection {detection_idx + 1}/{len(detection_results)}")
  1058. try:
  1059. #extracting
  1060. x_coords = [point[0] for point in detection_result]
  1061. y_coords = [point[1] for point in detection_result]
  1062. #min and max values for x and y
  1063. x_min, x_max = min(x_coords), max(x_coords)
  1064. y_min, y_max = min(y_coords), max(y_coords)
  1065. 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}")
  1066. except Exception as e:
  1067. logger.error(f"Error extracting bounding box coordinates for detection {detection_idx + 1}: {e}")
  1068. continue
  1069. # Converting bounding box values to integers for cropping
  1070. x_min, y_min, x_max, y_max = self.expand_bbox(x_min, y_min, x_max, y_max)
  1071. # Step 3: Cropping image based on detection output
  1072. try:
  1073. cropped_image = self.pixmap.toImage().copy(x_min, y_min, x_max - x_min, y_max - y_min)
  1074. cropped_image = labelme.utils.img_qt_to_arr(cropped_image)
  1075. orig_height, orig_width = cropped_image.shape[:2]
  1076. logger.debug(f"Original height: {orig_height}, Original width: {orig_width}")
  1077. # if the height is greater than the width we rotate for segmentaion
  1078. if orig_height > orig_width:
  1079. cropped_image = cv2.rotate(cropped_image, cv2.ROTATE_90_CLOCKWISE)
  1080. logger.debug(f"Rotated cropped image for detection {detection_idx + 1} due to height > width.")
  1081. orig_cropped_shape = cropped_image.shape[:2]
  1082. rotated = True
  1083. else:
  1084. rotated = False
  1085. # Save crop image
  1086. cv2.imwrite(f"cropped_image_{detection_idx + 1}.png", cropped_image)
  1087. logger.debug(f"Saved cropped image for detection {detection_idx + 1}: {cropped_image.shape}")
  1088. # logger.debug(f"Cropped image shape for detection {detection_idx + 1}: {cropped_image.shape}")
  1089. except Exception as e:
  1090. logger.error(f"Error cropping the image for detection {detection_idx + 1}: {e}")
  1091. continue
  1092. # Step 4: Resize the cropped image to match segmentation input size (1 64 256)
  1093. try:
  1094. orig_cropped_shape = cropped_image.shape[:2] # Save the original cropped image size
  1095. preprocessed_img = self._detection_model.preprocess_image(cropped_image, for_segmentation=True)
  1096. logger.debug(f"Preprocessed image shape for segmentation detection {detection_idx + 1}: {preprocessed_img.shape}")
  1097. except Exception as e:
  1098. logger.error(f"Error preprocessing the image for segmentation for detection {detection_idx + 1}: {e}")
  1099. continue
  1100. # Step 5: inference on segmentation model on cropped image
  1101. try:
  1102. seg_result = self._segmentation_model.segmentation_sess.infer_new_request({'x': preprocessed_img})
  1103. logger.debug(f"Segmentation model inference completed for detection {detection_idx + 1}.")
  1104. except Exception as e:
  1105. logger.error(f"Error during segmentation model inference for detection {detection_idx + 1}: {e}")
  1106. continue
  1107. # Step 6: Convert binary mask to polygon (contours)
  1108. try:
  1109. mask = seg_result['save_infer_model/scale_0.tmp_0'] #model output name
  1110. mask = mask.squeeze() # Remove batch dimension, should result in (64, 256)
  1111. logger.debug(f"Segmentation mask shape for detection {detection_idx + 1}: {mask.shape}")
  1112. # Normalize the mask to 0 and 255 and convert to uint8
  1113. mask = (mask * 255).astype(np.uint8)
  1114. # logger.debug(f"Converted mask shape for detection {detection_idx + 1}: {mask.shape}, dtype: {mask.dtype}")
  1115. # cv2.imwrite(f"segmentation_mask_{detection_idx + 1}.png", mask)
  1116. # Scale mask to original cropped image size
  1117. scaled_mask = cv2.resize(mask, (orig_cropped_shape[1], orig_cropped_shape[0]), interpolation=cv2.INTER_NEAREST)
  1118. if rotated:
  1119. cropped_image = cv2.rotate(cropped_image, cv2.ROTATE_90_COUNTERCLOCKWISE)
  1120. mask = cv2.rotate(scaled_mask, cv2.ROTATE_90_COUNTERCLOCKWISE)
  1121. rotated_cropped_shape = cropped_image.shape[:2]
  1122. else:
  1123. mask = scaled_mask
  1124. cv2.imwrite(f"scaled_segmentation_mask_{detection_idx + 1}.png", mask)
  1125. logger.debug(f"Saved segmentation mask for detection {detection_idx + 1}.")
  1126. # Step 7: Find contours
  1127. contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
  1128. logger.debug(f"Found {len(contours)} contours in the mask for detection {detection_idx + 1}.")
  1129. if len(contours) > 0:
  1130. # Get the largest contour
  1131. largest_contour = max(contours, key=cv2.contourArea)
  1132. # Get bounding and rotated rectangles
  1133. rot_rect = cv2.minAreaRect(largest_contour)
  1134. box_points = cv2.boxPoints(rot_rect)
  1135. box_points = np.int32(box_points)
  1136. # Log the rotated rectangle points
  1137. # logger.debug(f"Rotated corners (box_points) for detection {detection_idx + 1}: {box_points}")
  1138. # Step 1: Draw the rotated rectangle on the cropped image
  1139. cropped_with_rects = cropped_image.copy()
  1140. cv2.drawContours(cropped_with_rects, [box_points], -1, (0, 255, 0), 2)
  1141. cv2.imwrite(f"cropped_with_rects_{detection_idx + 1}.png", cropped_with_rects)
  1142. logger.debug(f"Saved cropped image with rectangles for detection {detection_idx + 1}.")
  1143. if rotated:
  1144. # Scale points from mask space to the original image space with rotation
  1145. scaled_points = self.scale_points(box_points, mask.shape, rotated_cropped_shape, x_min, y_min)
  1146. else:
  1147. # Scale points without rotation
  1148. scaled_points = self.scale_points(box_points, mask.shape, orig_cropped_shape, x_min, y_min)
  1149. self.pred_bbox_points = [scaled_points]
  1150. logger.debug(f"Predicted Bounding Box Points for barcode {detection_idx + 1}: {self.pred_bbox_points}")
  1151. if self.pred_bbox_points:
  1152. for bbox_points in self.pred_bbox_points:
  1153. try:
  1154. original_image = labelme.utils.img_qt_to_arr(self.pixmap.toImage())
  1155. # Call the decoding function
  1156. decoded_output = self._decoding_model.decode_from_points(
  1157. points=bbox_points,
  1158. detection_idx=detection_idx,
  1159. original_image = original_image
  1160. )
  1161. logger.debug(f"Decoded output for detection {detection_idx + 1}: {decoded_output}")
  1162. #passed to popup through app.py and label_dailog.py
  1163. self.decoded_barcode = decoded_output
  1164. # (Optional) Use the decoded output for updating UI or further processing
  1165. # Example: Print or log the decoded text
  1166. # print(f"Decoded output: {decoded_output}")
  1167. except Exception as e:
  1168. logger.error(f"Error during decoding for bbox points {bbox_points}: {e}")
  1169. ## pass each shape points to decode function here, do the decodign and give back the output here and print the output
  1170. if self.pred_bbox_points:
  1171. self.draw_pred = True
  1172. self.finalise()
  1173. else:
  1174. logger.info(f"No bounding boxes detected for detection {detection_idx + 1}.")
  1175. # Collect the segmentation result
  1176. all_segmentation_results.append(self.pred_bbox_points)
  1177. except Exception as e:
  1178. logger.error(f"Error creating the polygon shape for detection {detection_idx + 1}: {e}")
  1179. # **Reset critical variables after each detection**:
  1180. self.pred_bbox_points = None
  1181. self.draw_pred = False
  1182. # You now have a list of segmentation results for all detections
  1183. if all_segmentation_results:
  1184. logger.info(f"Segmentation results for all detections: {all_segmentation_results}")
  1185. def keyReleaseEvent(self, ev):
  1186. modifiers = ev.modifiers()
  1187. if self.drawing():
  1188. if int(modifiers) == 0:
  1189. self.snapping = True
  1190. elif self.editing():
  1191. if self.movingShape and self.selectedShapes:
  1192. index = self.shapes.index(self.selectedShapes[0])
  1193. if self.shapesBackups[-1][index].points != self.shapes[index].points:
  1194. self.storeShapes()
  1195. self.shapeMoved.emit()
  1196. self.movingShape = False
  1197. def setLastLabel(self, text, flags):
  1198. assert text
  1199. if(self.shapes is None or len(self.shapes)==0):
  1200. return
  1201. self.shapes[-1].label = text
  1202. self.shapes[-1].flags = flags
  1203. self.shapesBackups.pop()
  1204. self.storeShapes()
  1205. return self.shapes[-1]
  1206. def undoLastLine(self):
  1207. if(self.shapes is None or len(self.shapes)==0):
  1208. return
  1209. self.current = self.shapes.pop()
  1210. self.current.setOpen()
  1211. self.current.restoreShapeRaw()
  1212. if self.createMode in ["polygon", "linestrip"] and self.draw_pred is False:
  1213. self.line.points = [self.current[-1], self.current[0]]
  1214. elif self.createMode in ["rectangle", "line", "circle"]:
  1215. self.current.points = self.current.points[0:1]
  1216. elif self.createMode == "point":
  1217. self.current = None
  1218. self.drawingPolygon.emit(True)
  1219. def undoLastPoint(self):
  1220. if not self.current or self.current.isClosed():
  1221. return
  1222. self.current.popPoint()
  1223. if len(self.current) > 0:
  1224. self.line[0] = self.current[-1]
  1225. else:
  1226. self.current = None
  1227. self.drawingPolygon.emit(False)
  1228. self.update()
  1229. def loadPixmap(self, pixmap, clear_shapes=True):
  1230. self.pixmap = pixmap
  1231. if clear_shapes:
  1232. self.shapes = []
  1233. self.update()
  1234. def loadShapes(self, shapes, replace=True):
  1235. if replace:
  1236. self.shapes = list(shapes)
  1237. else:
  1238. self.shapes.extend(shapes)
  1239. self.storeShapes()
  1240. self.current = None
  1241. self.hShape = None
  1242. self.hVertex = None
  1243. self.hEdge = None
  1244. self.update()
  1245. def setShapeVisible(self, shape, value):
  1246. self.visible[shape] = value
  1247. self.update()
  1248. def overrideCursor(self, cursor):
  1249. self.restoreCursor()
  1250. self._cursor = cursor
  1251. QtWidgets.QApplication.setOverrideCursor(cursor)
  1252. def restoreCursor(self):
  1253. QtWidgets.QApplication.restoreOverrideCursor()
  1254. def resetState(self):
  1255. self.restoreCursor()
  1256. self.pixmap = None
  1257. self.shapesBackups = []
  1258. self.update()