canvas.py 38 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055
  1. import gdown
  2. from qtpy import QtCore
  3. from qtpy import QtGui
  4. from qtpy import QtWidgets
  5. import labelme.ai
  6. from labelme.logger import logger
  7. from labelme import QT5
  8. from labelme.shape import Shape
  9. import labelme.utils
  10. # TODO(unknown):
  11. # - [maybe] Find optimal epsilon value.
  12. CURSOR_DEFAULT = QtCore.Qt.ArrowCursor
  13. CURSOR_POINT = QtCore.Qt.PointingHandCursor
  14. CURSOR_DRAW = QtCore.Qt.CrossCursor
  15. CURSOR_MOVE = QtCore.Qt.ClosedHandCursor
  16. CURSOR_GRAB = QtCore.Qt.OpenHandCursor
  17. MOVE_SPEED = 5.0
  18. class Canvas(QtWidgets.QWidget):
  19. zoomRequest = QtCore.Signal(int, QtCore.QPoint)
  20. scrollRequest = QtCore.Signal(int, int)
  21. newShape = QtCore.Signal()
  22. selectionChanged = QtCore.Signal(list)
  23. shapeMoved = QtCore.Signal()
  24. drawingPolygon = QtCore.Signal(bool)
  25. vertexSelected = QtCore.Signal(bool)
  26. CREATE, EDIT = 0, 1
  27. # polygon, rectangle, line, or point
  28. _createMode = "polygon"
  29. _fill_drawing = False
  30. def __init__(self, *args, **kwargs):
  31. self.epsilon = kwargs.pop("epsilon", 10.0)
  32. self.double_click = kwargs.pop("double_click", "close")
  33. if self.double_click not in [None, "close"]:
  34. raise ValueError(
  35. "Unexpected value for double_click event: {}".format(
  36. self.double_click
  37. )
  38. )
  39. self.num_backups = kwargs.pop("num_backups", 10)
  40. self._crosshair = kwargs.pop(
  41. "crosshair",
  42. {
  43. "polygon": False,
  44. "rectangle": True,
  45. "circle": False,
  46. "line": False,
  47. "point": False,
  48. "linestrip": False,
  49. "ai_polygon": False,
  50. },
  51. )
  52. super(Canvas, self).__init__(*args, **kwargs)
  53. # Initialise local state.
  54. self.mode = self.EDIT
  55. self.shapes = []
  56. self.shapesBackups = []
  57. self.current = None
  58. self.selectedShapes = [] # save the selected shapes here
  59. self.selectedShapesCopy = []
  60. # self.line represents:
  61. # - createMode == 'polygon': edge from last point to current
  62. # - createMode == 'rectangle': diagonal line of the rectangle
  63. # - createMode == 'line': the line
  64. # - createMode == 'point': the point
  65. self.line = Shape()
  66. self.prevPoint = QtCore.QPoint()
  67. self.prevMovePoint = QtCore.QPoint()
  68. self.offsets = QtCore.QPoint(), QtCore.QPoint()
  69. self.scale = 1.0
  70. self.pixmap = QtGui.QPixmap()
  71. self.visible = {}
  72. self._hideBackround = False
  73. self.hideBackround = False
  74. self.hShape = None
  75. self.prevhShape = None
  76. self.hVertex = None
  77. self.prevhVertex = None
  78. self.hEdge = None
  79. self.prevhEdge = None
  80. self.movingShape = False
  81. self.snapping = True
  82. self.hShapeIsSelected = False
  83. self._painter = QtGui.QPainter()
  84. self._cursor = CURSOR_DEFAULT
  85. # Menus:
  86. # 0: right-click without selection and dragging of shapes
  87. # 1: right-click with selection and dragging of shapes
  88. self.menus = (QtWidgets.QMenu(), QtWidgets.QMenu())
  89. # Set widget options.
  90. self.setMouseTracking(True)
  91. self.setFocusPolicy(QtCore.Qt.WheelFocus)
  92. self._ai_model = None
  93. def fillDrawing(self):
  94. return self._fill_drawing
  95. def setFillDrawing(self, value):
  96. self._fill_drawing = value
  97. @property
  98. def createMode(self):
  99. return self._createMode
  100. @createMode.setter
  101. def createMode(self, value):
  102. if value not in [
  103. "polygon",
  104. "rectangle",
  105. "circle",
  106. "line",
  107. "point",
  108. "linestrip",
  109. "ai_polygon",
  110. ]:
  111. raise ValueError("Unsupported createMode: %s" % value)
  112. self._createMode = value
  113. def initializeAiModel(self, name):
  114. if name not in [model.name for model in labelme.ai.MODELS]:
  115. raise ValueError("Unsupported ai model: %s" % name)
  116. model = [model for model in labelme.ai.MODELS if model.name == name][0]
  117. if self._ai_model is not None and self._ai_model.name == model.name:
  118. logger.debug("AI model is already initialized: %r" % model.name)
  119. else:
  120. logger.debug("Initializing AI model: %r" % model.name)
  121. self._ai_model = labelme.ai.SegmentAnythingModel(
  122. name=model.name,
  123. encoder_path=gdown.cached_download(
  124. url=model.encoder_weight.url,
  125. md5=model.encoder_weight.md5,
  126. ),
  127. decoder_path=gdown.cached_download(
  128. url=model.decoder_weight.url,
  129. md5=model.decoder_weight.md5,
  130. ),
  131. )
  132. self._ai_model.set_image(
  133. image=labelme.utils.img_qt_to_arr(self.pixmap.toImage())
  134. )
  135. def storeShapes(self):
  136. shapesBackup = []
  137. for shape in self.shapes:
  138. shapesBackup.append(shape.copy())
  139. if len(self.shapesBackups) > self.num_backups:
  140. self.shapesBackups = self.shapesBackups[-self.num_backups - 1 :]
  141. self.shapesBackups.append(shapesBackup)
  142. @property
  143. def isShapeRestorable(self):
  144. # We save the state AFTER each edit (not before) so for an
  145. # edit to be undoable, we expect the CURRENT and the PREVIOUS state
  146. # to be in the undo stack.
  147. if len(self.shapesBackups) < 2:
  148. return False
  149. return True
  150. def restoreShape(self):
  151. # This does _part_ of the job of restoring shapes.
  152. # The complete process is also done in app.py::undoShapeEdit
  153. # and app.py::loadShapes and our own Canvas::loadShapes function.
  154. if not self.isShapeRestorable:
  155. return
  156. self.shapesBackups.pop() # latest
  157. # The application will eventually call Canvas.loadShapes which will
  158. # push this right back onto the stack.
  159. shapesBackup = self.shapesBackups.pop()
  160. self.shapes = shapesBackup
  161. self.selectedShapes = []
  162. for shape in self.shapes:
  163. shape.selected = False
  164. self.update()
  165. def enterEvent(self, ev):
  166. self.overrideCursor(self._cursor)
  167. def leaveEvent(self, ev):
  168. self.unHighlight()
  169. self.restoreCursor()
  170. def focusOutEvent(self, ev):
  171. self.restoreCursor()
  172. def isVisible(self, shape):
  173. return self.visible.get(shape, True)
  174. def drawing(self):
  175. return self.mode == self.CREATE
  176. def editing(self):
  177. return self.mode == self.EDIT
  178. def setEditing(self, value=True):
  179. self.mode = self.EDIT if value else self.CREATE
  180. if self.mode == self.EDIT:
  181. # CREATE -> EDIT
  182. self.repaint() # clear crosshair
  183. else:
  184. # EDIT -> CREATE
  185. self.unHighlight()
  186. self.deSelectShape()
  187. def unHighlight(self):
  188. if self.hShape:
  189. self.hShape.highlightClear()
  190. self.update()
  191. self.prevhShape = self.hShape
  192. self.prevhVertex = self.hVertex
  193. self.prevhEdge = self.hEdge
  194. self.hShape = self.hVertex = self.hEdge = None
  195. def selectedVertex(self):
  196. return self.hVertex is not None
  197. def selectedEdge(self):
  198. return self.hEdge is not None
  199. def mouseMoveEvent(self, ev):
  200. """Update line with last point and current coordinates."""
  201. try:
  202. if QT5:
  203. pos = self.transformPos(ev.localPos())
  204. else:
  205. pos = self.transformPos(ev.posF())
  206. except AttributeError:
  207. return
  208. self.prevMovePoint = pos
  209. self.restoreCursor()
  210. is_shift_pressed = ev.modifiers() & QtCore.Qt.ShiftModifier
  211. # Polygon drawing.
  212. if self.drawing():
  213. if self.createMode == "ai_polygon":
  214. self.line.shape_type = "points"
  215. else:
  216. self.line.shape_type = self.createMode
  217. self.overrideCursor(CURSOR_DRAW)
  218. if not self.current:
  219. self.repaint() # draw crosshair
  220. return
  221. if self.outOfPixmap(pos):
  222. # Don't allow the user to draw outside the pixmap.
  223. # Project the point to the pixmap's edges.
  224. pos = self.intersectionPoint(self.current[-1], pos)
  225. elif (
  226. self.snapping
  227. and len(self.current) > 1
  228. and self.createMode == "polygon"
  229. and self.closeEnough(pos, self.current[0])
  230. ):
  231. # Attract line to starting point and
  232. # colorise to alert the user.
  233. pos = self.current[0]
  234. self.overrideCursor(CURSOR_POINT)
  235. self.current.highlightVertex(0, Shape.NEAR_VERTEX)
  236. if self.createMode in ["polygon", "linestrip"]:
  237. self.line.points = [self.current[-1], pos]
  238. self.line.point_labels = [1, 1]
  239. elif self.createMode == "ai_polygon":
  240. self.line.points = [self.current.points[-1], pos]
  241. self.line.point_labels = [
  242. self.current.point_labels[-1],
  243. 0 if is_shift_pressed else 1,
  244. ]
  245. elif self.createMode == "rectangle":
  246. self.line.points = [self.current[0], pos]
  247. self.line.point_labels = [1, 1]
  248. self.line.close()
  249. elif self.createMode == "circle":
  250. self.line.points = [self.current[0], pos]
  251. self.line.point_labels = [1, 1]
  252. self.line.shape_type = "circle"
  253. elif self.createMode == "line":
  254. self.line.points = [self.current[0], pos]
  255. self.line.point_labels = [1, 1]
  256. self.line.close()
  257. elif self.createMode == "point":
  258. self.line.points = [self.current[0]]
  259. self.line.point_labels = [1]
  260. self.line.close()
  261. assert len(self.line.points) == len(self.line.point_labels)
  262. self.repaint()
  263. self.current.highlightClear()
  264. return
  265. # Polygon copy moving.
  266. if QtCore.Qt.RightButton & ev.buttons():
  267. if self.selectedShapesCopy and self.prevPoint:
  268. self.overrideCursor(CURSOR_MOVE)
  269. self.boundedMoveShapes(self.selectedShapesCopy, pos)
  270. self.repaint()
  271. elif self.selectedShapes:
  272. self.selectedShapesCopy = [
  273. s.copy() for s in self.selectedShapes
  274. ]
  275. self.repaint()
  276. return
  277. # Polygon/Vertex moving.
  278. if QtCore.Qt.LeftButton & ev.buttons():
  279. if self.selectedVertex():
  280. self.boundedMoveVertex(pos)
  281. self.repaint()
  282. self.movingShape = True
  283. elif self.selectedShapes and self.prevPoint:
  284. self.overrideCursor(CURSOR_MOVE)
  285. self.boundedMoveShapes(self.selectedShapes, pos)
  286. self.repaint()
  287. self.movingShape = True
  288. return
  289. # Just hovering over the canvas, 2 possibilities:
  290. # - Highlight shapes
  291. # - Highlight vertex
  292. # Update shape/vertex fill and tooltip value accordingly.
  293. self.setToolTip(self.tr("Image"))
  294. for shape in reversed([s for s in self.shapes if self.isVisible(s)]):
  295. # Look for a nearby vertex to highlight. If that fails,
  296. # check if we happen to be inside a shape.
  297. index = shape.nearestVertex(pos, self.epsilon / self.scale)
  298. index_edge = shape.nearestEdge(pos, self.epsilon / self.scale)
  299. if index is not None:
  300. if self.selectedVertex():
  301. self.hShape.highlightClear()
  302. self.prevhVertex = self.hVertex = index
  303. self.prevhShape = self.hShape = shape
  304. self.prevhEdge = self.hEdge
  305. self.hEdge = None
  306. shape.highlightVertex(index, shape.MOVE_VERTEX)
  307. self.overrideCursor(CURSOR_POINT)
  308. self.setToolTip(self.tr("Click & drag to move point"))
  309. self.setStatusTip(self.toolTip())
  310. self.update()
  311. break
  312. elif index_edge is not None and shape.canAddPoint():
  313. if self.selectedVertex():
  314. self.hShape.highlightClear()
  315. self.prevhVertex = self.hVertex
  316. self.hVertex = None
  317. self.prevhShape = self.hShape = shape
  318. self.prevhEdge = self.hEdge = index_edge
  319. self.overrideCursor(CURSOR_POINT)
  320. self.setToolTip(self.tr("Click to create point"))
  321. self.setStatusTip(self.toolTip())
  322. self.update()
  323. break
  324. elif shape.containsPoint(pos):
  325. if self.selectedVertex():
  326. self.hShape.highlightClear()
  327. self.prevhVertex = self.hVertex
  328. self.hVertex = None
  329. self.prevhShape = self.hShape = shape
  330. self.prevhEdge = self.hEdge
  331. self.hEdge = None
  332. self.setToolTip(
  333. self.tr("Click & drag to move shape '%s'") % shape.label
  334. )
  335. self.setStatusTip(self.toolTip())
  336. self.overrideCursor(CURSOR_GRAB)
  337. self.update()
  338. break
  339. else: # Nothing found, clear highlights, reset state.
  340. self.unHighlight()
  341. self.vertexSelected.emit(self.hVertex is not None)
  342. def addPointToEdge(self):
  343. shape = self.prevhShape
  344. index = self.prevhEdge
  345. point = self.prevMovePoint
  346. if shape is None or index is None or point is None:
  347. return
  348. shape.insertPoint(index, point)
  349. shape.highlightVertex(index, shape.MOVE_VERTEX)
  350. self.hShape = shape
  351. self.hVertex = index
  352. self.hEdge = None
  353. self.movingShape = True
  354. def removeSelectedPoint(self):
  355. shape = self.prevhShape
  356. index = self.prevhVertex
  357. if shape is None or index is None:
  358. return
  359. shape.removePoint(index)
  360. shape.highlightClear()
  361. self.hShape = shape
  362. self.prevhVertex = None
  363. self.movingShape = True # Save changes
  364. def mousePressEvent(self, ev):
  365. if QT5:
  366. pos = self.transformPos(ev.localPos())
  367. else:
  368. pos = self.transformPos(ev.posF())
  369. is_shift_pressed = ev.modifiers() & QtCore.Qt.ShiftModifier
  370. if ev.button() == QtCore.Qt.LeftButton:
  371. if self.drawing():
  372. if self.current:
  373. # Add point to existing shape.
  374. if self.createMode == "polygon":
  375. self.current.addPoint(self.line[1])
  376. self.line[0] = self.current[-1]
  377. if self.current.isClosed():
  378. self.finalise()
  379. elif self.createMode in ["rectangle", "circle", "line"]:
  380. assert len(self.current.points) == 1
  381. self.current.points = self.line.points
  382. self.finalise()
  383. elif self.createMode == "linestrip":
  384. self.current.addPoint(self.line[1])
  385. self.line[0] = self.current[-1]
  386. if int(ev.modifiers()) == QtCore.Qt.ControlModifier:
  387. self.finalise()
  388. elif self.createMode == "ai_polygon":
  389. self.current.addPoint(
  390. self.line.points[1],
  391. label=self.line.point_labels[1],
  392. )
  393. self.line.points[0] = self.current.points[-1]
  394. self.line.point_labels[0] = self.current.point_labels[
  395. -1
  396. ]
  397. if ev.modifiers() & QtCore.Qt.ControlModifier:
  398. self.finalise()
  399. elif not self.outOfPixmap(pos):
  400. # Create new shape.
  401. self.current = Shape(
  402. shape_type="points"
  403. if self.createMode == "ai_polygon"
  404. else self.createMode
  405. )
  406. self.current.addPoint(
  407. pos, label=0 if is_shift_pressed else 1
  408. )
  409. if self.createMode == "point":
  410. self.finalise()
  411. elif (
  412. self.createMode == "ai_polygon"
  413. and ev.modifiers() & QtCore.Qt.ControlModifier
  414. ):
  415. self.finalise()
  416. else:
  417. if self.createMode == "circle":
  418. self.current.shape_type = "circle"
  419. self.line.points = [pos, pos]
  420. if (
  421. self.createMode == "ai_polygon"
  422. and is_shift_pressed
  423. ):
  424. self.line.point_labels = [0, 0]
  425. else:
  426. self.line.point_labels = [1, 1]
  427. self.setHiding()
  428. self.drawingPolygon.emit(True)
  429. self.update()
  430. elif self.editing():
  431. if self.selectedEdge():
  432. self.addPointToEdge()
  433. elif (
  434. self.selectedVertex()
  435. and int(ev.modifiers()) == QtCore.Qt.ShiftModifier
  436. ):
  437. # Delete point if: left-click + SHIFT on a point
  438. self.removeSelectedPoint()
  439. group_mode = int(ev.modifiers()) == QtCore.Qt.ControlModifier
  440. self.selectShapePoint(pos, multiple_selection_mode=group_mode)
  441. self.prevPoint = pos
  442. self.repaint()
  443. elif ev.button() == QtCore.Qt.RightButton and self.editing():
  444. group_mode = int(ev.modifiers()) == QtCore.Qt.ControlModifier
  445. if not self.selectedShapes or (
  446. self.hShape is not None
  447. and self.hShape not in self.selectedShapes
  448. ):
  449. self.selectShapePoint(pos, multiple_selection_mode=group_mode)
  450. self.repaint()
  451. self.prevPoint = pos
  452. def mouseReleaseEvent(self, ev):
  453. if ev.button() == QtCore.Qt.RightButton:
  454. menu = self.menus[len(self.selectedShapesCopy) > 0]
  455. self.restoreCursor()
  456. if (
  457. not menu.exec_(self.mapToGlobal(ev.pos()))
  458. and self.selectedShapesCopy
  459. ):
  460. # Cancel the move by deleting the shadow copy.
  461. self.selectedShapesCopy = []
  462. self.repaint()
  463. elif ev.button() == QtCore.Qt.LeftButton:
  464. if self.editing():
  465. if (
  466. self.hShape is not None
  467. and self.hShapeIsSelected
  468. and not self.movingShape
  469. ):
  470. self.selectionChanged.emit(
  471. [x for x in self.selectedShapes if x != self.hShape]
  472. )
  473. if self.movingShape and self.hShape:
  474. index = self.shapes.index(self.hShape)
  475. if (
  476. self.shapesBackups[-1][index].points
  477. != self.shapes[index].points
  478. ):
  479. self.storeShapes()
  480. self.shapeMoved.emit()
  481. self.movingShape = False
  482. def endMove(self, copy):
  483. assert self.selectedShapes and self.selectedShapesCopy
  484. assert len(self.selectedShapesCopy) == len(self.selectedShapes)
  485. if copy:
  486. for i, shape in enumerate(self.selectedShapesCopy):
  487. self.shapes.append(shape)
  488. self.selectedShapes[i].selected = False
  489. self.selectedShapes[i] = shape
  490. else:
  491. for i, shape in enumerate(self.selectedShapesCopy):
  492. self.selectedShapes[i].points = shape.points
  493. self.selectedShapesCopy = []
  494. self.repaint()
  495. self.storeShapes()
  496. return True
  497. def hideBackroundShapes(self, value):
  498. self.hideBackround = value
  499. if self.selectedShapes:
  500. # Only hide other shapes if there is a current selection.
  501. # Otherwise the user will not be able to select a shape.
  502. self.setHiding(True)
  503. self.update()
  504. def setHiding(self, enable=True):
  505. self._hideBackround = self.hideBackround if enable else False
  506. def canCloseShape(self):
  507. return self.drawing() and self.current and len(self.current) > 2
  508. def mouseDoubleClickEvent(self, ev):
  509. if self.double_click != "close":
  510. return
  511. if (
  512. self.createMode == "polygon" and self.canCloseShape()
  513. ) or self.createMode == "ai_polygon":
  514. self.finalise()
  515. def selectShapes(self, shapes):
  516. self.setHiding()
  517. self.selectionChanged.emit(shapes)
  518. self.update()
  519. def selectShapePoint(self, point, multiple_selection_mode):
  520. """Select the first shape created which contains this point."""
  521. if self.selectedVertex(): # A vertex is marked for selection.
  522. index, shape = self.hVertex, self.hShape
  523. shape.highlightVertex(index, shape.MOVE_VERTEX)
  524. else:
  525. for shape in reversed(self.shapes):
  526. if self.isVisible(shape) and shape.containsPoint(point):
  527. self.setHiding()
  528. if shape not in self.selectedShapes:
  529. if multiple_selection_mode:
  530. self.selectionChanged.emit(
  531. self.selectedShapes + [shape]
  532. )
  533. else:
  534. self.selectionChanged.emit([shape])
  535. self.hShapeIsSelected = False
  536. else:
  537. self.hShapeIsSelected = True
  538. self.calculateOffsets(point)
  539. return
  540. self.deSelectShape()
  541. def calculateOffsets(self, point):
  542. left = self.pixmap.width() - 1
  543. right = 0
  544. top = self.pixmap.height() - 1
  545. bottom = 0
  546. for s in self.selectedShapes:
  547. rect = s.boundingRect()
  548. if rect.left() < left:
  549. left = rect.left()
  550. if rect.right() > right:
  551. right = rect.right()
  552. if rect.top() < top:
  553. top = rect.top()
  554. if rect.bottom() > bottom:
  555. bottom = rect.bottom()
  556. x1 = left - point.x()
  557. y1 = top - point.y()
  558. x2 = right - point.x()
  559. y2 = bottom - point.y()
  560. self.offsets = QtCore.QPointF(x1, y1), QtCore.QPointF(x2, y2)
  561. def boundedMoveVertex(self, pos):
  562. index, shape = self.hVertex, self.hShape
  563. point = shape[index]
  564. if self.outOfPixmap(pos):
  565. pos = self.intersectionPoint(point, pos)
  566. shape.moveVertexBy(index, pos - point)
  567. def boundedMoveShapes(self, shapes, pos):
  568. if self.outOfPixmap(pos):
  569. return False # No need to move
  570. o1 = pos + self.offsets[0]
  571. if self.outOfPixmap(o1):
  572. pos -= QtCore.QPointF(min(0, o1.x()), min(0, o1.y()))
  573. o2 = pos + self.offsets[1]
  574. if self.outOfPixmap(o2):
  575. pos += QtCore.QPointF(
  576. min(0, self.pixmap.width() - o2.x()),
  577. min(0, self.pixmap.height() - o2.y()),
  578. )
  579. # XXX: The next line tracks the new position of the cursor
  580. # relative to the shape, but also results in making it
  581. # a bit "shaky" when nearing the border and allows it to
  582. # go outside of the shape's area for some reason.
  583. # self.calculateOffsets(self.selectedShapes, pos)
  584. dp = pos - self.prevPoint
  585. if dp:
  586. for shape in shapes:
  587. shape.moveBy(dp)
  588. self.prevPoint = pos
  589. return True
  590. return False
  591. def deSelectShape(self):
  592. if self.selectedShapes:
  593. self.setHiding(False)
  594. self.selectionChanged.emit([])
  595. self.hShapeIsSelected = False
  596. self.update()
  597. def deleteSelected(self):
  598. deleted_shapes = []
  599. if self.selectedShapes:
  600. for shape in self.selectedShapes:
  601. self.shapes.remove(shape)
  602. deleted_shapes.append(shape)
  603. self.storeShapes()
  604. self.selectedShapes = []
  605. self.update()
  606. return deleted_shapes
  607. def deleteShape(self, shape):
  608. if shape in self.selectedShapes:
  609. self.selectedShapes.remove(shape)
  610. if shape in self.shapes:
  611. self.shapes.remove(shape)
  612. self.storeShapes()
  613. self.update()
  614. def duplicateSelectedShapes(self):
  615. if self.selectedShapes:
  616. self.selectedShapesCopy = [s.copy() for s in self.selectedShapes]
  617. self.boundedShiftShapes(self.selectedShapesCopy)
  618. self.endMove(copy=True)
  619. return self.selectedShapes
  620. def boundedShiftShapes(self, shapes):
  621. # Try to move in one direction, and if it fails in another.
  622. # Give up if both fail.
  623. point = shapes[0][0]
  624. offset = QtCore.QPointF(2.0, 2.0)
  625. self.offsets = QtCore.QPoint(), QtCore.QPoint()
  626. self.prevPoint = point
  627. if not self.boundedMoveShapes(shapes, point - offset):
  628. self.boundedMoveShapes(shapes, point + offset)
  629. def paintEvent(self, event):
  630. if not self.pixmap:
  631. return super(Canvas, self).paintEvent(event)
  632. p = self._painter
  633. p.begin(self)
  634. p.setRenderHint(QtGui.QPainter.Antialiasing)
  635. p.setRenderHint(QtGui.QPainter.HighQualityAntialiasing)
  636. p.setRenderHint(QtGui.QPainter.SmoothPixmapTransform)
  637. p.scale(self.scale, self.scale)
  638. p.translate(self.offsetToCenter())
  639. p.drawPixmap(0, 0, self.pixmap)
  640. # draw crosshair
  641. if (
  642. self._crosshair[self._createMode]
  643. and self.drawing()
  644. and self.prevMovePoint
  645. and not self.outOfPixmap(self.prevMovePoint)
  646. ):
  647. p.setPen(QtGui.QColor(0, 0, 0))
  648. p.drawLine(
  649. 0,
  650. int(self.prevMovePoint.y()),
  651. self.width() - 1,
  652. int(self.prevMovePoint.y()),
  653. )
  654. p.drawLine(
  655. int(self.prevMovePoint.x()),
  656. 0,
  657. int(self.prevMovePoint.x()),
  658. self.height() - 1,
  659. )
  660. Shape.scale = self.scale
  661. for shape in self.shapes:
  662. if (shape.selected or not self._hideBackround) and self.isVisible(
  663. shape
  664. ):
  665. shape.fill = shape.selected or shape == self.hShape
  666. shape.paint(p)
  667. if self.current:
  668. self.current.paint(p)
  669. assert len(self.line.points) == len(self.line.point_labels)
  670. self.line.paint(p)
  671. if self.selectedShapesCopy:
  672. for s in self.selectedShapesCopy:
  673. s.paint(p)
  674. if (
  675. self.fillDrawing()
  676. and self.createMode == "polygon"
  677. and self.current is not None
  678. and len(self.current.points) >= 2
  679. ):
  680. drawing_shape = self.current.copy()
  681. if drawing_shape.fill_color.getRgb()[3] == 0:
  682. logger.warning(
  683. "fill_drawing=true, but fill_color is transparent,"
  684. " so forcing to be opaque."
  685. )
  686. drawing_shape.fill_color.setAlpha(64)
  687. drawing_shape.addPoint(self.line[1])
  688. drawing_shape.fill = True
  689. drawing_shape.paint(p)
  690. elif self.createMode == "ai_polygon" and self.current is not None:
  691. drawing_shape = self.current.copy()
  692. drawing_shape.addPoint(
  693. point=self.line.points[1],
  694. label=self.line.point_labels[1],
  695. )
  696. drawing_shape.selected = True
  697. points = self._ai_model.predict_polygon_from_points(
  698. points=[
  699. [point.x(), point.y()] for point in drawing_shape.points
  700. ],
  701. point_labels=drawing_shape.point_labels,
  702. )
  703. if len(points) > 2:
  704. drawing_shape.setShapeRefined(
  705. points=[
  706. QtCore.QPointF(point[0], point[1]) for point in points
  707. ],
  708. point_labels=[1] * len(points),
  709. shape_type="polygon",
  710. )
  711. drawing_shape.fill = self.fillDrawing()
  712. drawing_shape.paint(p)
  713. p.end()
  714. def transformPos(self, point):
  715. """Convert from widget-logical coordinates to painter-logical ones."""
  716. return point / self.scale - self.offsetToCenter()
  717. def offsetToCenter(self):
  718. s = self.scale
  719. area = super(Canvas, self).size()
  720. w, h = self.pixmap.width() * s, self.pixmap.height() * s
  721. aw, ah = area.width(), area.height()
  722. x = (aw - w) / (2 * s) if aw > w else 0
  723. y = (ah - h) / (2 * s) if ah > h else 0
  724. return QtCore.QPointF(x, y)
  725. def outOfPixmap(self, p):
  726. w, h = self.pixmap.width(), self.pixmap.height()
  727. return not (0 <= p.x() <= w - 1 and 0 <= p.y() <= h - 1)
  728. def finalise(self):
  729. assert self.current
  730. if self.createMode == "ai_polygon":
  731. # convert points to polygon by an AI model
  732. assert self.current.shape_type == "points"
  733. points = self._ai_model.predict_polygon_from_points(
  734. points=[
  735. [point.x(), point.y()] for point in self.current.points
  736. ],
  737. point_labels=self.current.point_labels,
  738. )
  739. self.current.setShapeRefined(
  740. points=[
  741. QtCore.QPointF(point[0], point[1]) for point in points
  742. ],
  743. point_labels=[1] * len(points),
  744. shape_type="polygon",
  745. )
  746. self.current.close()
  747. self.shapes.append(self.current)
  748. self.storeShapes()
  749. self.current = None
  750. self.setHiding(False)
  751. self.newShape.emit()
  752. self.update()
  753. def closeEnough(self, p1, p2):
  754. # d = distance(p1 - p2)
  755. # m = (p1-p2).manhattanLength()
  756. # print "d %.2f, m %d, %.2f" % (d, m, d - m)
  757. # divide by scale to allow more precision when zoomed in
  758. return labelme.utils.distance(p1 - p2) < (self.epsilon / self.scale)
  759. def intersectionPoint(self, p1, p2):
  760. # Cycle through each image edge in clockwise fashion,
  761. # and find the one intersecting the current line segment.
  762. # http://paulbourke.net/geometry/lineline2d/
  763. size = self.pixmap.size()
  764. points = [
  765. (0, 0),
  766. (size.width() - 1, 0),
  767. (size.width() - 1, size.height() - 1),
  768. (0, size.height() - 1),
  769. ]
  770. # x1, y1 should be in the pixmap, x2, y2 should be out of the pixmap
  771. x1 = min(max(p1.x(), 0), size.width() - 1)
  772. y1 = min(max(p1.y(), 0), size.height() - 1)
  773. x2, y2 = p2.x(), p2.y()
  774. d, i, (x, y) = min(self.intersectingEdges((x1, y1), (x2, y2), points))
  775. x3, y3 = points[i]
  776. x4, y4 = points[(i + 1) % 4]
  777. if (x, y) == (x1, y1):
  778. # Handle cases where previous point is on one of the edges.
  779. if x3 == x4:
  780. return QtCore.QPointF(x3, min(max(0, y2), max(y3, y4)))
  781. else: # y3 == y4
  782. return QtCore.QPointF(min(max(0, x2), max(x3, x4)), y3)
  783. return QtCore.QPointF(x, y)
  784. def intersectingEdges(self, point1, point2, points):
  785. """Find intersecting edges.
  786. For each edge formed by `points', yield the intersection
  787. with the line segment `(x1,y1) - (x2,y2)`, if it exists.
  788. Also return the distance of `(x2,y2)' to the middle of the
  789. edge along with its index, so that the one closest can be chosen.
  790. """
  791. (x1, y1) = point1
  792. (x2, y2) = point2
  793. for i in range(4):
  794. x3, y3 = points[i]
  795. x4, y4 = points[(i + 1) % 4]
  796. denom = (y4 - y3) * (x2 - x1) - (x4 - x3) * (y2 - y1)
  797. nua = (x4 - x3) * (y1 - y3) - (y4 - y3) * (x1 - x3)
  798. nub = (x2 - x1) * (y1 - y3) - (y2 - y1) * (x1 - x3)
  799. if denom == 0:
  800. # This covers two cases:
  801. # nua == nub == 0: Coincident
  802. # otherwise: Parallel
  803. continue
  804. ua, ub = nua / denom, nub / denom
  805. if 0 <= ua <= 1 and 0 <= ub <= 1:
  806. x = x1 + ua * (x2 - x1)
  807. y = y1 + ua * (y2 - y1)
  808. m = QtCore.QPointF((x3 + x4) / 2, (y3 + y4) / 2)
  809. d = labelme.utils.distance(m - QtCore.QPointF(x2, y2))
  810. yield d, i, (x, y)
  811. # These two, along with a call to adjustSize are required for the
  812. # scroll area.
  813. def sizeHint(self):
  814. return self.minimumSizeHint()
  815. def minimumSizeHint(self):
  816. if self.pixmap:
  817. return self.scale * self.pixmap.size()
  818. return super(Canvas, self).minimumSizeHint()
  819. def wheelEvent(self, ev):
  820. if QT5:
  821. mods = ev.modifiers()
  822. delta = ev.angleDelta()
  823. if QtCore.Qt.ControlModifier == int(mods):
  824. # with Ctrl/Command key
  825. # zoom
  826. self.zoomRequest.emit(delta.y(), ev.pos())
  827. else:
  828. # scroll
  829. self.scrollRequest.emit(delta.x(), QtCore.Qt.Horizontal)
  830. self.scrollRequest.emit(delta.y(), QtCore.Qt.Vertical)
  831. else:
  832. if ev.orientation() == QtCore.Qt.Vertical:
  833. mods = ev.modifiers()
  834. if QtCore.Qt.ControlModifier == int(mods):
  835. # with Ctrl/Command key
  836. self.zoomRequest.emit(ev.delta(), ev.pos())
  837. else:
  838. self.scrollRequest.emit(
  839. ev.delta(),
  840. QtCore.Qt.Horizontal
  841. if (QtCore.Qt.ShiftModifier == int(mods))
  842. else QtCore.Qt.Vertical,
  843. )
  844. else:
  845. self.scrollRequest.emit(ev.delta(), QtCore.Qt.Horizontal)
  846. ev.accept()
  847. def moveByKeyboard(self, offset):
  848. if self.selectedShapes:
  849. self.boundedMoveShapes(
  850. self.selectedShapes, self.prevPoint + offset
  851. )
  852. self.repaint()
  853. self.movingShape = True
  854. def keyPressEvent(self, ev):
  855. modifiers = ev.modifiers()
  856. key = ev.key()
  857. if self.drawing():
  858. if key == QtCore.Qt.Key_Escape and self.current:
  859. self.current = None
  860. self.drawingPolygon.emit(False)
  861. self.update()
  862. elif key == QtCore.Qt.Key_Return and self.canCloseShape():
  863. self.finalise()
  864. elif modifiers == QtCore.Qt.AltModifier:
  865. self.snapping = False
  866. elif self.editing():
  867. if key == QtCore.Qt.Key_Up:
  868. self.moveByKeyboard(QtCore.QPointF(0.0, -MOVE_SPEED))
  869. elif key == QtCore.Qt.Key_Down:
  870. self.moveByKeyboard(QtCore.QPointF(0.0, MOVE_SPEED))
  871. elif key == QtCore.Qt.Key_Left:
  872. self.moveByKeyboard(QtCore.QPointF(-MOVE_SPEED, 0.0))
  873. elif key == QtCore.Qt.Key_Right:
  874. self.moveByKeyboard(QtCore.QPointF(MOVE_SPEED, 0.0))
  875. def keyReleaseEvent(self, ev):
  876. modifiers = ev.modifiers()
  877. if self.drawing():
  878. if int(modifiers) == 0:
  879. self.snapping = True
  880. elif self.editing():
  881. if self.movingShape and self.selectedShapes:
  882. index = self.shapes.index(self.selectedShapes[0])
  883. if (
  884. self.shapesBackups[-1][index].points
  885. != self.shapes[index].points
  886. ):
  887. self.storeShapes()
  888. self.shapeMoved.emit()
  889. self.movingShape = False
  890. def setLastLabel(self, text, flags):
  891. assert text
  892. self.shapes[-1].label = text
  893. self.shapes[-1].flags = flags
  894. self.shapesBackups.pop()
  895. self.storeShapes()
  896. return self.shapes[-1]
  897. def undoLastLine(self):
  898. assert self.shapes
  899. self.current = self.shapes.pop()
  900. self.current.setOpen()
  901. self.current.restoreShapeRaw()
  902. if self.createMode in ["polygon", "linestrip"]:
  903. self.line.points = [self.current[-1], self.current[0]]
  904. elif self.createMode in ["rectangle", "line", "circle"]:
  905. self.current.points = self.current.points[0:1]
  906. elif self.createMode == "point":
  907. self.current = None
  908. self.drawingPolygon.emit(True)
  909. def undoLastPoint(self):
  910. if not self.current or self.current.isClosed():
  911. return
  912. self.current.popPoint()
  913. if len(self.current) > 0:
  914. self.line[0] = self.current[-1]
  915. else:
  916. self.current = None
  917. self.drawingPolygon.emit(False)
  918. self.update()
  919. def loadPixmap(self, pixmap, clear_shapes=True):
  920. self.pixmap = pixmap
  921. if self._ai_model:
  922. self._ai_model.set_image(
  923. image=labelme.utils.img_qt_to_arr(self.pixmap.toImage())
  924. )
  925. if clear_shapes:
  926. self.shapes = []
  927. self.update()
  928. def loadShapes(self, shapes, replace=True):
  929. if replace:
  930. self.shapes = list(shapes)
  931. else:
  932. self.shapes.extend(shapes)
  933. self.storeShapes()
  934. self.current = None
  935. self.hShape = None
  936. self.hVertex = None
  937. self.hEdge = None
  938. self.update()
  939. def setShapeVisible(self, shape, value):
  940. self.visible[shape] = value
  941. self.update()
  942. def overrideCursor(self, cursor):
  943. self.restoreCursor()
  944. self._cursor = cursor
  945. QtWidgets.QApplication.setOverrideCursor(cursor)
  946. def restoreCursor(self):
  947. QtWidgets.QApplication.restoreOverrideCursor()
  948. def resetState(self):
  949. self.restoreCursor()
  950. self.pixmap = None
  951. self.shapesBackups = []
  952. self.update()