PPOCRLabel.py 79 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747174817491750175117521753175417551756175717581759176017611762176317641765176617671768176917701771177217731774177517761777177817791780178117821783178417851786178717881789179017911792179317941795179617971798179918001801180218031804180518061807180818091810181118121813181418151816181718181819182018211822182318241825182618271828182918301831183218331834183518361837183818391840184118421843184418451846184718481849185018511852185318541855185618571858185918601861186218631864186518661867186818691870187118721873187418751876187718781879188018811882188318841885188618871888188918901891189218931894189518961897189818991900190119021903190419051906190719081909191019111912191319141915191619171918191919201921192219231924192519261927192819291930193119321933193419351936193719381939194019411942194319441945194619471948194919501951195219531954195519561957195819591960196119621963196419651966196719681969197019711972197319741975197619771978197919801981198219831984198519861987198819891990199119921993199419951996199719981999200020012002200320042005
  1. # Copyright (c) <2015-Present> Tzutalin
  2. # Copyright (C) 2013 MIT, Computer Science and Artificial Intelligence Laboratory. Bryan Russell, Antonio Torralba,
  3. # William T. Freeman. Permission is hereby granted, free of charge, to any person obtaining a copy of this software and
  4. # associated documentation files (the "Software"), to deal in the Software without restriction, including without
  5. # limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the
  6. # Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
  7. # The above copyright notice and this permission notice shall be included in all copies or substantial portions of
  8. # the Software. THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT
  9. # NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
  10. # SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF
  11. # CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
  12. # THE SOFTWARE.
  13. #!/usr/bin/env python
  14. # -*- coding: utf-8 -*-
  15. # pyrcc5 -o libs/resources.py resources.qrc
  16. import argparse
  17. import ast
  18. import codecs
  19. import os.path
  20. import platform
  21. import subprocess
  22. import sys
  23. from functools import partial
  24. from collections import defaultdict
  25. import json
  26. import cv2
  27. __dir__ = os.path.dirname(os.path.abspath(__file__))
  28. sys.path.append(__dir__)
  29. sys.path.append(os.path.abspath(os.path.join(__dir__, '../..')))
  30. sys.path.append("..")
  31. from paddleocr import PaddleOCR
  32. try:
  33. from PyQt5 import QtCore, QtGui, QtWidgets
  34. from PyQt5.QtGui import *
  35. from PyQt5.QtCore import *
  36. from PyQt5.QtWidgets import *
  37. except ImportError:
  38. # needed for py3+qt4
  39. # Ref:
  40. # http://pyqt.sourceforge.net/Docs/PyQt4/incompatible_apis.html
  41. # http://stackoverflow.com/questions/21217399/pyqt4-qtcore-qvariant-object-instead-of-a-string
  42. if sys.version_info.major >= 3:
  43. import sip
  44. sip.setapi('QVariant', 2)
  45. from PyQt4.QtGui import *
  46. from PyQt4.QtCore import *
  47. from combobox import ComboBox
  48. from libs.constants import *
  49. from libs.utils import *
  50. from libs.settings import Settings
  51. from libs.shape import Shape, DEFAULT_LINE_COLOR, DEFAULT_FILL_COLOR
  52. from libs.stringBundle import StringBundle
  53. from libs.canvas import Canvas
  54. from libs.zoomWidget import ZoomWidget
  55. from libs.autoDialog import AutoDialog
  56. from libs.labelDialog import LabelDialog
  57. from libs.colorDialog import ColorDialog
  58. from libs.toolBar import ToolBar
  59. from libs.ustr import ustr
  60. from libs.hashableQListWidgetItem import HashableQListWidgetItem
  61. from libs.editinlist import EditInList
  62. __appname__ = 'PPOCRLabel'
  63. class WindowMixin(object):
  64. def menu(self, title, actions=None):
  65. menu = self.menuBar().addMenu(title)
  66. if actions:
  67. addActions(menu, actions)
  68. return menu
  69. def toolbar(self, title, actions=None):
  70. toolbar = ToolBar(title)
  71. toolbar.setObjectName(u'%sToolBar' % title)
  72. # toolbar.setOrientation(Qt.Vertical)
  73. toolbar.setToolButtonStyle(Qt.ToolButtonTextUnderIcon)
  74. if actions:
  75. addActions(toolbar, actions)
  76. self.addToolBar(Qt.LeftToolBarArea, toolbar)
  77. return toolbar
  78. class MainWindow(QMainWindow, WindowMixin):
  79. FIT_WINDOW, FIT_WIDTH, MANUAL_ZOOM = list(range(3))
  80. def __init__(self, lang="ch", defaultFilename=None, defaultPrefdefClassFile=None, defaultSaveDir=None):
  81. super(MainWindow, self).__init__()
  82. self.setWindowTitle(__appname__)
  83. # Load setting in the main thread
  84. self.settings = Settings()
  85. self.settings.load()
  86. settings = self.settings
  87. self.lang = lang
  88. # Load string bundle for i18n
  89. if lang not in ['ch', 'en']:
  90. lang = 'en'
  91. self.stringBundle = StringBundle.getBundle(localeStr='zh-CN' if lang=='ch' else 'en') # 'en'
  92. getStr = lambda strId: self.stringBundle.getString(strId)
  93. self.defaultSaveDir = defaultSaveDir
  94. self.ocr = PaddleOCR(use_pdserving=False, use_angle_cls=True, det=True, cls=True, use_gpu=False, lang=lang)
  95. if os.path.exists('./data/paddle.png'):
  96. result = self.ocr.ocr('./data/paddle.png', cls=True, det=True)
  97. # For loading all image under a directory
  98. self.mImgList = []
  99. self.mImgList5 = []
  100. self.dirname = None
  101. self.labelHist = []
  102. self.lastOpenDir = None
  103. self.result_dic = []
  104. self.changeFileFolder = False
  105. self.haveAutoReced = False
  106. self.labelFile = None
  107. self.currIndex = 0
  108. # Whether we need to save or not.
  109. self.dirty = False
  110. self._noSelectionSlot = False
  111. self._beginner = True
  112. self.screencastViewer = self.getAvailableScreencastViewer()
  113. self.screencast = "https://github.com/PaddlePaddle/PaddleOCR"
  114. # Load predefined classes to the list
  115. self.loadPredefinedClasses(defaultPrefdefClassFile)
  116. # Main widgets and related state.
  117. self.labelDialog = LabelDialog(parent=self, listItem=self.labelHist)
  118. self.autoDialog = AutoDialog(parent=self)
  119. self.itemsToShapes = {}
  120. self.shapesToItems = {}
  121. self.itemsToShapesbox = {}
  122. self.shapesToItemsbox = {}
  123. self.prevLabelText = getStr('tempLabel')
  124. self.model = 'paddle'
  125. self.PPreader = None
  126. self.autoSaveNum = 5
  127. ################# file list ###############
  128. self.fileListWidget = QListWidget()
  129. self.fileListWidget.itemClicked.connect(self.fileitemDoubleClicked)
  130. self.fileListWidget.setIconSize(QSize(25, 25))
  131. filelistLayout = QVBoxLayout()
  132. filelistLayout.setContentsMargins(0, 0, 0, 0)
  133. filelistLayout.addWidget(self.fileListWidget)
  134. self.AutoRecognition = QToolButton()
  135. self.AutoRecognition.setToolButtonStyle(Qt.ToolButtonTextBesideIcon)
  136. self.AutoRecognition.setIcon(newIcon('Auto'))
  137. autoRecLayout = QHBoxLayout()
  138. autoRecLayout.setContentsMargins(0, 0, 0, 0)
  139. autoRecLayout.addWidget(self.AutoRecognition)
  140. autoRecContainer = QWidget()
  141. autoRecContainer.setLayout(autoRecLayout)
  142. filelistLayout.addWidget(autoRecContainer)
  143. fileListContainer = QWidget()
  144. fileListContainer.setLayout(filelistLayout)
  145. self.filedock = QDockWidget(getStr('fileList'), self)
  146. self.filedock.setObjectName(getStr('files'))
  147. self.filedock.setWidget(fileListContainer)
  148. self.addDockWidget(Qt.LeftDockWidgetArea, self.filedock)
  149. ######## Right area ##########
  150. listLayout = QVBoxLayout()
  151. listLayout.setContentsMargins(0, 0, 0, 0)
  152. self.editButton = QToolButton()
  153. self.reRecogButton = QToolButton()
  154. self.reRecogButton.setIcon(newIcon('reRec', 30))
  155. self.reRecogButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon)
  156. self.newButton = QToolButton()
  157. self.newButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon)
  158. self.SaveButton = QToolButton()
  159. self.SaveButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon)
  160. self.DelButton = QToolButton()
  161. self.DelButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon)
  162. lefttoptoolbox = QHBoxLayout()
  163. lefttoptoolbox.addWidget(self.newButton)
  164. lefttoptoolbox.addWidget(self.reRecogButton)
  165. lefttoptoolboxcontainer = QWidget()
  166. lefttoptoolboxcontainer.setLayout(lefttoptoolbox)
  167. listLayout.addWidget(lefttoptoolboxcontainer)
  168. ################## label list ####################
  169. # Create and add a widget for showing current label items
  170. self.labelList = EditInList()
  171. labelListContainer = QWidget()
  172. labelListContainer.setLayout(listLayout)
  173. self.labelList.itemActivated.connect(self.labelSelectionChanged)
  174. self.labelList.itemSelectionChanged.connect(self.labelSelectionChanged)
  175. self.labelList.clicked.connect(self.labelList.item_clicked)
  176. # Connect to itemChanged to detect checkbox changes.
  177. self.labelList.itemChanged.connect(self.labelItemChanged)
  178. self.labelListDock = QDockWidget(getStr('recognitionResult'),self)
  179. self.labelListDock.setWidget(self.labelList)
  180. self.labelListDock.setFeatures(QDockWidget.NoDockWidgetFeatures)
  181. listLayout.addWidget(self.labelListDock)
  182. ################## detection box ####################
  183. self.BoxList = QListWidget()
  184. self.BoxList.itemActivated.connect(self.boxSelectionChanged)
  185. self.BoxList.itemSelectionChanged.connect(self.boxSelectionChanged)
  186. self.BoxList.itemDoubleClicked.connect(self.editBox)
  187. # Connect to itemChanged to detect checkbox changes.
  188. self.BoxList.itemChanged.connect(self.boxItemChanged)
  189. self.BoxListDock = QDockWidget(getStr('detectionBoxposition'), self)
  190. self.BoxListDock.setWidget(self.BoxList)
  191. self.BoxListDock.setFeatures(QDockWidget.NoDockWidgetFeatures)
  192. listLayout.addWidget(self.BoxListDock)
  193. ############ lower right area ############
  194. leftbtmtoolbox = QHBoxLayout()
  195. leftbtmtoolbox.addWidget(self.SaveButton)
  196. leftbtmtoolbox.addWidget(self.DelButton)
  197. leftbtmtoolboxcontainer = QWidget()
  198. leftbtmtoolboxcontainer.setLayout(leftbtmtoolbox)
  199. listLayout.addWidget(leftbtmtoolboxcontainer)
  200. self.dock = QDockWidget(getStr('boxLabelText'), self)
  201. self.dock.setObjectName(getStr('labels'))
  202. self.dock.setWidget(labelListContainer)
  203. ########## zoom bar #########
  204. self.imgsplider = QSlider(Qt.Horizontal)
  205. self.imgsplider.valueChanged.connect(self.CanvasSizeChange)
  206. self.imgsplider.setMinimum(-150)
  207. self.imgsplider.setMaximum(150)
  208. self.imgsplider.setSingleStep(1)
  209. self.imgsplider.setTickPosition(QSlider.TicksBelow)
  210. self.imgsplider.setTickInterval(1)
  211. op = QGraphicsOpacityEffect()
  212. op.setOpacity(0.2)
  213. self.imgsplider.setGraphicsEffect(op)
  214. # self.imgsplider.setAttribute(Qt.WA_TranslucentBackground)
  215. self.imgsplider.setStyleSheet("background-color:transparent")
  216. self.imgsliderDock = QDockWidget(getStr('ImageResize'), self)
  217. self.imgsliderDock.setObjectName(getStr('IR'))
  218. self.imgsliderDock.setWidget(self.imgsplider)
  219. self.imgsliderDock.setFeatures(QDockWidget.DockWidgetFloatable)
  220. self.imgsliderDock.setAttribute(Qt.WA_TranslucentBackground)
  221. self.addDockWidget(Qt.RightDockWidgetArea, self.imgsliderDock)
  222. self.zoomWidget = ZoomWidget()
  223. self.colorDialog = ColorDialog(parent=self)
  224. self.zoomWidgetValue = self.zoomWidget.value()
  225. ########## thumbnail #########
  226. hlayout = QHBoxLayout()
  227. m = (0, 0, 0, 0)
  228. hlayout.setSpacing(0)
  229. hlayout.setContentsMargins(*m)
  230. self.preButton = QToolButton()
  231. self.preButton.setIcon(newIcon("prev",40))
  232. self.preButton.setIconSize(QSize(40, 100))
  233. self.preButton.clicked.connect(self.openPrevImg)
  234. self.preButton.setStyleSheet('border: none;')
  235. self.preButton.setShortcut('a')
  236. self.iconlist = QListWidget()
  237. self.iconlist.setViewMode(QListView.IconMode)
  238. self.iconlist.setFlow(QListView.TopToBottom)
  239. self.iconlist.setSpacing(10)
  240. self.iconlist.setIconSize(QSize(50, 50))
  241. self.iconlist.setMovement(False)
  242. self.iconlist.setResizeMode(QListView.Adjust)
  243. self.iconlist.itemClicked.connect(self.iconitemDoubleClicked)
  244. self.iconlist.setStyleSheet("background-color:transparent; border: none;")
  245. self.iconlist.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff)
  246. self.nextButton = QToolButton()
  247. self.nextButton.setIcon(newIcon("next", 40))
  248. self.nextButton.setIconSize(QSize(40, 100))
  249. self.nextButton.setStyleSheet('border: none;')
  250. self.nextButton.clicked.connect(self.openNextImg)
  251. self.nextButton.setShortcut('d')
  252. hlayout.addWidget(self.preButton)
  253. hlayout.addWidget(self.iconlist)
  254. hlayout.addWidget(self.nextButton)
  255. iconListContainer = QWidget()
  256. iconListContainer.setLayout(hlayout)
  257. iconListContainer.setFixedHeight(100)
  258. ########### Canvas ###########
  259. self.canvas = Canvas(parent=self)
  260. self.canvas.zoomRequest.connect(self.zoomRequest)
  261. self.canvas.setDrawingShapeToSquare(settings.get(SETTING_DRAW_SQUARE, False))
  262. scroll = QScrollArea()
  263. scroll.setWidget(self.canvas)
  264. scroll.setWidgetResizable(True)
  265. self.scrollBars = {
  266. Qt.Vertical: scroll.verticalScrollBar(),
  267. Qt.Horizontal: scroll.horizontalScrollBar()
  268. }
  269. self.scrollArea = scroll
  270. self.canvas.scrollRequest.connect(self.scrollRequest)
  271. self.canvas.newShape.connect(partial(self.newShape, False))
  272. self.canvas.shapeMoved.connect(self.updateBoxlist) # self.setDirty
  273. self.canvas.selectionChanged.connect(self.shapeSelectionChanged)
  274. self.canvas.drawingPolygon.connect(self.toggleDrawingSensitive)
  275. centerLayout = QVBoxLayout()
  276. centerLayout.setContentsMargins(0, 0, 0, 0)
  277. centerLayout.addWidget(scroll)
  278. #centerLayout.addWidget(self.icondock)
  279. centerLayout.addWidget(iconListContainer,0,Qt.AlignCenter)
  280. centercontainer = QWidget()
  281. centercontainer.setLayout(centerLayout)
  282. # self.scrolldock = QDockWidget('WorkSpace',self)
  283. # self.scrolldock.setObjectName('WorkSpace')
  284. # self.scrolldock.setWidget(centercontainer)
  285. # self.scrolldock.setFeatures(QDockWidget.NoDockWidgetFeatures)
  286. # orititle = self.scrolldock.titleBarWidget()
  287. # tmpwidget = QWidget()
  288. # self.scrolldock.setTitleBarWidget(tmpwidget)
  289. # del orititle
  290. self.setCentralWidget(centercontainer) #self.scrolldock
  291. self.addDockWidget(Qt.RightDockWidgetArea, self.dock)
  292. # self.filedock.setFeatures(QDockWidget.DockWidgetFloatable)
  293. self.filedock.setFeatures(self.filedock.features() ^ QDockWidget.DockWidgetFloatable)
  294. self.dockFeatures = QDockWidget.DockWidgetClosable | QDockWidget.DockWidgetFloatable
  295. self.dock.setFeatures(self.dock.features() ^ self.dockFeatures)
  296. self.filedock.setFeatures(QDockWidget.NoDockWidgetFeatures)
  297. ###### Actions #######
  298. action = partial(newAction, self)
  299. quit = action(getStr('quit'), self.close,
  300. 'Ctrl+Q', 'quit', getStr('quitApp'))
  301. opendir = action(getStr('openDir'), self.openDirDialog,
  302. 'Ctrl+u', 'open', getStr('openDir'))
  303. save = action(getStr('save'), self.saveFile,
  304. 'Ctrl+V', 'verify', getStr('saveDetail'), enabled=False)
  305. alcm = action(getStr('choosemodel'), self.autolcm,
  306. 'Ctrl+M', 'next', getStr('tipchoosemodel'))
  307. deleteImg = action(getStr('deleteImg'), self.deleteImg, 'Ctrl+Shift+D', 'close', getStr('deleteImgDetail'),
  308. enabled=True)
  309. resetAll = action(getStr('resetAll'), self.resetAll, None, 'resetall', getStr('resetAllDetail'))
  310. color1 = action(getStr('boxLineColor'), self.chooseColor1,
  311. 'Ctrl+L', 'color_line', getStr('boxLineColorDetail'))
  312. createMode = action(getStr('crtBox'), self.setCreateMode,
  313. 'w', 'new', getStr('crtBoxDetail'), enabled=False)
  314. editMode = action('&Edit\nRectBox', self.setEditMode,
  315. 'Ctrl+J', 'edit', u'Move and edit Boxs', enabled=False)
  316. create = action(getStr('crtBox'), self.createShape,
  317. 'w', 'new', getStr('crtBoxDetail'), enabled=False)
  318. delete = action(getStr('delBox'), self.deleteSelectedShape,
  319. 'backspace', 'delete', getStr('delBoxDetail'), enabled=False)
  320. copy = action(getStr('dupBox'), self.copySelectedShape,
  321. 'Ctrl+C', 'copy', getStr('dupBoxDetail'),
  322. enabled=False)
  323. hideAll = action(getStr('hideBox'), partial(self.togglePolygons, False),
  324. 'Ctrl+H', 'hide', getStr('hideAllBoxDetail'),
  325. enabled=False)
  326. showAll = action(getStr('showBox'), partial(self.togglePolygons, True),
  327. 'Ctrl+A', 'hide', getStr('showAllBoxDetail'),
  328. enabled=False)
  329. help = action(getStr('tutorial'), self.showTutorialDialog, None, 'help', getStr('tutorialDetail'))
  330. showInfo = action(getStr('info'), self.showInfoDialog, None, 'help', getStr('info'))
  331. showSteps = action(getStr('steps'), self.showStepsDialog, None, 'help', getStr('steps'))
  332. zoom = QWidgetAction(self)
  333. zoom.setDefaultWidget(self.zoomWidget)
  334. self.zoomWidget.setWhatsThis(
  335. u"Zoom in or out of the image. Also accessible with"
  336. " %s and %s from the canvas." % (fmtShortcut("Ctrl+[-+]"),
  337. fmtShortcut("Ctrl+Wheel")))
  338. self.zoomWidget.setEnabled(False)
  339. zoomIn = action(getStr('zoomin'), partial(self.addZoom, 10),
  340. 'Ctrl++', 'zoom-in', getStr('zoominDetail'), enabled=False)
  341. zoomOut = action(getStr('zoomout'), partial(self.addZoom, -10),
  342. 'Ctrl+-', 'zoom-out', getStr('zoomoutDetail'), enabled=False)
  343. zoomOrg = action(getStr('originalsize'), partial(self.setZoom, 100),
  344. 'Ctrl+=', 'zoom', getStr('originalsizeDetail'), enabled=False)
  345. fitWindow = action(getStr('fitWin'), self.setFitWindow,
  346. 'Ctrl+F', 'fit-window', getStr('fitWinDetail'),
  347. checkable=True, enabled=False)
  348. fitWidth = action(getStr('fitWidth'), self.setFitWidth,
  349. 'Ctrl+Shift+F', 'fit-width', getStr('fitWidthDetail'),
  350. checkable=True, enabled=False)
  351. # Group zoom controls into a list for easier toggling.
  352. zoomActions = (self.zoomWidget, zoomIn, zoomOut,
  353. zoomOrg, fitWindow, fitWidth)
  354. self.zoomMode = self.MANUAL_ZOOM
  355. self.scalers = {
  356. self.FIT_WINDOW: self.scaleFitWindow,
  357. self.FIT_WIDTH: self.scaleFitWidth,
  358. # Set to one to scale to 100% when loading files.
  359. self.MANUAL_ZOOM: lambda: 1,
  360. }
  361. edit = action(getStr('editLabel'), self.editLabel,
  362. 'Ctrl+E', 'edit', getStr('editLabelDetail'),
  363. enabled=False)
  364. ######## New actions #######
  365. AutoRec = action(getStr('autoRecognition'), self.autoRecognition,
  366. 'Ctrl+Shift+A', 'Auto', getStr('autoRecognition'), enabled=False)
  367. reRec = action(getStr('reRecognition'), self.reRecognition,
  368. 'Ctrl+Shift+R', 'reRec', getStr('reRecognition'), enabled=False)
  369. singleRere = action(getStr('singleRe'), self.singleRerecognition,
  370. 'Ctrl+R', 'reRec', getStr('singleRe'), enabled=False)
  371. createpoly = action(getStr('creatPolygon'), self.createPolygon,
  372. 'q', 'new', 'Creat Polygon', enabled=True)
  373. saveRec = action(getStr('saveRec'), self.saveRecResult,
  374. '', 'save', getStr('saveRec'), enabled=False)
  375. saveLabel = action(getStr('saveLabel'), self.saveLabelFile, #
  376. 'Ctrl+S', 'save', getStr('saveLabel'), enabled=False)
  377. self.editButton.setDefaultAction(edit)
  378. self.newButton.setDefaultAction(create)
  379. self.DelButton.setDefaultAction(deleteImg)
  380. self.SaveButton.setDefaultAction(save)
  381. self.AutoRecognition.setDefaultAction(AutoRec)
  382. self.reRecogButton.setDefaultAction(reRec)
  383. # self.preButton.setDefaultAction(openPrevImg)
  384. # self.nextButton.setDefaultAction(openNextImg)
  385. ############# Zoom layout ##############
  386. zoomLayout = QHBoxLayout()
  387. zoomLayout.addStretch()
  388. self.zoominButton = QToolButton()
  389. self.zoominButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon)
  390. self.zoominButton.setDefaultAction(zoomIn)
  391. self.zoomoutButton = QToolButton()
  392. self.zoomoutButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon)
  393. self.zoomoutButton.setDefaultAction(zoomOut)
  394. self.zoomorgButton = QToolButton()
  395. self.zoomorgButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon)
  396. self.zoomorgButton.setDefaultAction(zoomOrg)
  397. zoomLayout.addWidget(self.zoominButton)
  398. zoomLayout.addWidget(self.zoomorgButton)
  399. zoomLayout.addWidget(self.zoomoutButton)
  400. zoomContainer = QWidget()
  401. zoomContainer.setLayout(zoomLayout)
  402. zoomContainer.setGeometry(0, 0, 30, 150)
  403. shapeLineColor = action(getStr('shapeLineColor'), self.chshapeLineColor,
  404. icon='color_line', tip=getStr('shapeLineColorDetail'),
  405. enabled=False)
  406. shapeFillColor = action(getStr('shapeFillColor'), self.chshapeFillColor,
  407. icon='color', tip=getStr('shapeFillColorDetail'),
  408. enabled=False)
  409. # Label list context menu.
  410. labelMenu = QMenu()
  411. addActions(labelMenu, (edit, delete))
  412. self.labelList.setContextMenuPolicy(Qt.CustomContextMenu)
  413. self.labelList.customContextMenuRequested.connect(
  414. self.popLabelListMenu)
  415. # Draw squares/rectangles
  416. self.drawSquaresOption = QAction(getStr('drawSquares'), self)
  417. self.drawSquaresOption.setCheckable(True)
  418. self.drawSquaresOption.setChecked(settings.get(SETTING_DRAW_SQUARE, False))
  419. self.drawSquaresOption.triggered.connect(self.toogleDrawSquare)
  420. # Store actions for further handling.
  421. self.actions = struct(save=save, resetAll=resetAll, deleteImg=deleteImg,
  422. lineColor=color1, create=create, delete=delete, edit=edit, copy=copy,
  423. saveRec=saveRec, singleRere=singleRere,AutoRec=AutoRec,reRec=reRec,
  424. createMode=createMode, editMode=editMode,
  425. shapeLineColor=shapeLineColor, shapeFillColor=shapeFillColor,
  426. zoom=zoom, zoomIn=zoomIn, zoomOut=zoomOut, zoomOrg=zoomOrg,
  427. fitWindow=fitWindow, fitWidth=fitWidth,
  428. zoomActions=zoomActions, saveLabel=saveLabel,
  429. fileMenuActions=(
  430. opendir, saveLabel, resetAll, quit),
  431. beginner=(), advanced=(),
  432. editMenu=(createpoly, edit, copy, delete,singleRere,
  433. None, color1, self.drawSquaresOption),
  434. beginnerContext=(create, edit, copy, delete, singleRere),
  435. advancedContext=(createMode, editMode, edit, copy,
  436. delete, shapeLineColor, shapeFillColor),
  437. onLoadActive=(
  438. create, createMode, editMode),
  439. onShapesPresent=(hideAll, showAll))
  440. # menus
  441. self.menus = struct(
  442. file=self.menu('&'+getStr('mfile')),
  443. edit=self.menu('&'+getStr('medit')),
  444. view=self.menu('&'+getStr('mview')),
  445. autolabel=self.menu('&PaddleOCR'),
  446. help=self.menu('&'+getStr('mhelp')),
  447. recentFiles=QMenu('Open &Recent'),
  448. labelList=labelMenu)
  449. self.lastLabel = None
  450. # Add option to enable/disable labels being displayed at the top of bounding boxes
  451. self.displayLabelOption = QAction(getStr('displayLabel'), self)
  452. self.displayLabelOption.setShortcut("Ctrl+Shift+P")
  453. self.displayLabelOption.setCheckable(True)
  454. self.displayLabelOption.setChecked(settings.get(SETTING_PAINT_LABEL, False))
  455. self.displayLabelOption.triggered.connect(self.togglePaintLabelsOption)
  456. self.labelDialogOption = QAction(getStr('labelDialogOption'), self)
  457. self.labelDialogOption.setShortcut("Ctrl+Shift+L")
  458. self.labelDialogOption.setCheckable(True)
  459. self.labelDialogOption.setChecked(settings.get(SETTING_PAINT_LABEL, False))
  460. self.labelDialogOption.triggered.connect(self.speedChoose)
  461. addActions(self.menus.file,
  462. (opendir, None, saveLabel, saveRec, None, resetAll, deleteImg, quit))
  463. addActions(self.menus.help, (showSteps, showInfo))
  464. addActions(self.menus.view, (
  465. self.displayLabelOption, self.labelDialogOption,
  466. None,
  467. hideAll, showAll, None,
  468. zoomIn, zoomOut, zoomOrg, None,
  469. fitWindow, fitWidth))
  470. addActions(self.menus.autolabel, (AutoRec, reRec, alcm, None, help)) #
  471. self.menus.file.aboutToShow.connect(self.updateFileMenu)
  472. # Custom context menu for the canvas widget:
  473. addActions(self.canvas.menus[0], self.actions.beginnerContext)
  474. addActions(self.canvas.menus[1], (
  475. action('&Copy here', self.copyShape),
  476. action('&Move here', self.moveShape)))
  477. self.statusBar().showMessage('%s started.' % __appname__)
  478. self.statusBar().show()
  479. # Application state.
  480. self.image = QImage()
  481. self.filePath = ustr(defaultFilename)
  482. self.lastOpenDir = None
  483. self.recentFiles = []
  484. self.maxRecent = 7
  485. self.lineColor = None
  486. self.fillColor = None
  487. self.zoom_level = 100
  488. self.fit_window = False
  489. # Add Chris
  490. self.difficult = False
  491. ## Fix the compatible issue for qt4 and qt5. Convert the QStringList to python list
  492. if settings.get(SETTING_RECENT_FILES):
  493. if have_qstring():
  494. recentFileQStringList = settings.get(SETTING_RECENT_FILES)
  495. self.recentFiles = [ustr(i) for i in recentFileQStringList]
  496. else:
  497. self.recentFiles = recentFileQStringList = settings.get(SETTING_RECENT_FILES)
  498. size = settings.get(SETTING_WIN_SIZE, QSize(1200, 800))
  499. position = QPoint(0, 0)
  500. saved_position = settings.get(SETTING_WIN_POSE, position)
  501. # Fix the multiple monitors issue
  502. for i in range(QApplication.desktop().screenCount()):
  503. if QApplication.desktop().availableGeometry(i).contains(saved_position):
  504. position = saved_position
  505. break
  506. self.resize(size)
  507. self.move(position)
  508. saveDir = ustr(settings.get(SETTING_SAVE_DIR, None))
  509. self.lastOpenDir = ustr(settings.get(SETTING_LAST_OPEN_DIR, None))
  510. self.restoreState(settings.get(SETTING_WIN_STATE, QByteArray()))
  511. Shape.line_color = self.lineColor = QColor(settings.get(SETTING_LINE_COLOR, DEFAULT_LINE_COLOR))
  512. Shape.fill_color = self.fillColor = QColor(settings.get(SETTING_FILL_COLOR, DEFAULT_FILL_COLOR))
  513. self.canvas.setDrawingColor(self.lineColor)
  514. # Add chris
  515. Shape.difficult = self.difficult
  516. # ADD:
  517. # Populate the File menu dynamically.
  518. self.updateFileMenu()
  519. # Since loading the file may take some time, make sure it runs in the background.
  520. if self.filePath and os.path.isdir(self.filePath):
  521. self.queueEvent(partial(self.importDirImages, self.filePath or ""))
  522. elif self.filePath:
  523. self.queueEvent(partial(self.loadFile, self.filePath or ""))
  524. # Callbacks:
  525. self.zoomWidget.valueChanged.connect(self.paintCanvas)
  526. self.populateModeActions()
  527. # Display cursor coordinates at the right of status bar
  528. self.labelCoordinates = QLabel('')
  529. self.statusBar().addPermanentWidget(self.labelCoordinates)
  530. # Open Dir if deafult file
  531. if self.filePath and os.path.isdir(self.filePath):
  532. self.openDirDialog(dirpath=self.filePath, silent=True)
  533. def keyReleaseEvent(self, event):
  534. if event.key() == Qt.Key_Control:
  535. self.canvas.setDrawingShapeToSquare(False)
  536. def keyPressEvent(self, event):
  537. if event.key() == Qt.Key_Control:
  538. # Draw rectangle if Ctrl is pressed
  539. self.canvas.setDrawingShapeToSquare(True)
  540. def noShapes(self):
  541. return not self.itemsToShapes
  542. def populateModeActions(self):
  543. self.canvas.menus[0].clear()
  544. addActions(self.canvas.menus[0], self.actions.beginnerContext)
  545. self.menus.edit.clear()
  546. actions = (self.actions.create,) # if self.beginner() else (self.actions.createMode, self.actions.editMode)
  547. addActions(self.menus.edit, actions + self.actions.editMenu)
  548. def setDirty(self):
  549. self.dirty = True
  550. self.actions.save.setEnabled(True)
  551. def setClean(self):
  552. self.dirty = False
  553. self.actions.save.setEnabled(False)
  554. self.actions.create.setEnabled(True)
  555. def toggleActions(self, value=True):
  556. """Enable/Disable widgets which depend on an opened image."""
  557. for z in self.actions.zoomActions:
  558. z.setEnabled(value)
  559. for action in self.actions.onLoadActive:
  560. action.setEnabled(value)
  561. def queueEvent(self, function):
  562. QTimer.singleShot(0, function)
  563. def status(self, message, delay=5000):
  564. self.statusBar().showMessage(message, delay)
  565. def resetState(self):
  566. self.itemsToShapes.clear()
  567. self.shapesToItems.clear()
  568. self.itemsToShapesbox.clear() # ADD
  569. self.shapesToItemsbox.clear()
  570. self.labelList.clear()
  571. self.BoxList.clear()
  572. self.filePath = None
  573. self.imageData = None
  574. self.labelFile = None
  575. self.canvas.resetState()
  576. self.labelCoordinates.clear()
  577. # self.comboBox.cb.clear()
  578. self.result_dic = []
  579. def currentItem(self):
  580. items = self.labelList.selectedItems()
  581. if items:
  582. return items[0]
  583. return None
  584. def currentBox(self):
  585. items = self.BoxList.selectedItems()
  586. if items:
  587. return items[0]
  588. return None
  589. def addRecentFile(self, filePath):
  590. if filePath in self.recentFiles:
  591. self.recentFiles.remove(filePath)
  592. elif len(self.recentFiles) >= self.maxRecent:
  593. self.recentFiles.pop()
  594. self.recentFiles.insert(0, filePath)
  595. def beginner(self):
  596. return self._beginner
  597. def advanced(self):
  598. return not self.beginner()
  599. def getAvailableScreencastViewer(self):
  600. osName = platform.system()
  601. if osName == 'Windows':
  602. return ['C:\\Program Files\\Internet Explorer\\iexplore.exe']
  603. elif osName == 'Linux':
  604. return ['xdg-open']
  605. elif osName == 'Darwin':
  606. return ['open']
  607. ## Callbacks ##
  608. def showTutorialDialog(self):
  609. subprocess.Popen(self.screencastViewer + [self.screencast])
  610. def showInfoDialog(self):
  611. from libs.__init__ import __version__
  612. msg = u'Name:{0} \nApp Version:{1} \n{2} '.format(__appname__, __version__, sys.version_info)
  613. QMessageBox.information(self, u'Information', msg)
  614. def showStepsDialog(self):
  615. msg = stepsInfo(self.lang)
  616. QMessageBox.information(self, u'Information', msg)
  617. def createShape(self):
  618. assert self.beginner()
  619. self.canvas.setEditing(False)
  620. self.actions.create.setEnabled(False)
  621. self.canvas.fourpoint = False
  622. def createPolygon(self):
  623. assert self.beginner()
  624. self.canvas.setEditing(False)
  625. self.canvas.fourpoint = True
  626. self.actions.create.setEnabled(False)
  627. def toggleDrawingSensitive(self, drawing=True):
  628. """In the middle of drawing, toggling between modes should be disabled."""
  629. self.actions.editMode.setEnabled(not drawing)
  630. if not drawing and self.beginner():
  631. # Cancel creation.
  632. print('Cancel creation.')
  633. self.canvas.setEditing(True)
  634. self.canvas.restoreCursor()
  635. self.actions.create.setEnabled(True)
  636. def toggleDrawMode(self, edit=True):
  637. self.canvas.setEditing(edit)
  638. self.actions.createMode.setEnabled(edit)
  639. self.actions.editMode.setEnabled(not edit)
  640. def setCreateMode(self):
  641. assert self.advanced()
  642. self.toggleDrawMode(False)
  643. def setEditMode(self):
  644. assert self.advanced()
  645. self.toggleDrawMode(True)
  646. self.labelSelectionChanged()
  647. def updateFileMenu(self):
  648. currFilePath = self.filePath
  649. def exists(filename):
  650. return os.path.exists(filename)
  651. menu = self.menus.recentFiles
  652. menu.clear()
  653. files = [f for f in self.recentFiles if f !=
  654. currFilePath and exists(f)]
  655. for i, f in enumerate(files):
  656. icon = newIcon('labels')
  657. action = QAction(
  658. icon, '&%d %s' % (i + 1, QFileInfo(f).fileName()), self)
  659. action.triggered.connect(partial(self.loadRecent, f))
  660. menu.addAction(action)
  661. def popLabelListMenu(self, point):
  662. self.menus.labelList.exec_(self.labelList.mapToGlobal(point))
  663. def editLabel(self):
  664. if not self.canvas.editing():
  665. return
  666. item = self.currentItem()
  667. if not item:
  668. return
  669. text = self.labelDialog.popUp(item.text())
  670. if text is not None:
  671. item.setText(text)
  672. # item.setBackground(generateColorByText(text))
  673. self.setDirty()
  674. self.updateComboBox()
  675. ######## detection box related functions #######
  676. def boxItemChanged(self, item):
  677. shape = self.itemsToShapesbox[item]
  678. box = ast.literal_eval(item.text())
  679. # print('shape in labelItemChanged is',shape.points)
  680. if box != [(p.x(), p.y()) for p in shape.points]:
  681. # shape.points = box
  682. shape.points = [QPointF(p[0], p[1]) for p in box]
  683. # QPointF(x,y)
  684. # shape.line_color = generateColorByText(shape.label)
  685. self.setDirty()
  686. else: # User probably changed item visibility
  687. self.canvas.setShapeVisible(shape, True)#item.checkState() == Qt.Checked
  688. def editBox(self): # ADD
  689. if not self.canvas.editing():
  690. return
  691. item = self.currentBox()
  692. if not item:
  693. return
  694. text = self.labelDialog.popUp(item.text())
  695. imageSize = str(self.image.size())
  696. width, height = self.image.width(), self.image.height()
  697. if text:
  698. try:
  699. text_list = eval(text)
  700. except:
  701. msg_box = QMessageBox(QMessageBox.Warning, 'Warning', 'Please enter the correct format')
  702. msg_box.exec_()
  703. return
  704. if len(text_list) < 4:
  705. msg_box = QMessageBox(QMessageBox.Warning, 'Warning', 'Please enter the coordinates of 4 points')
  706. msg_box.exec_()
  707. return
  708. for box in text_list:
  709. if box[0] > width or box[0] < 0 or box[1] > height or box[1] < 0:
  710. msg_box = QMessageBox(QMessageBox.Warning, 'Warning', 'Out of picture size')
  711. msg_box.exec_()
  712. return
  713. item.setText(text)
  714. # item.setBackground(generateColorByText(text))
  715. self.setDirty()
  716. self.updateComboBox()
  717. def updateBoxlist(self):
  718. shape = self.canvas.selectedShape
  719. item = self.shapesToItemsbox[shape] # listitem
  720. text = [(int(p.x()), int(p.y())) for p in shape.points]
  721. item.setText(str(text))
  722. self.setDirty()
  723. def indexTo5Files(self, currIndex):
  724. if currIndex < 2:
  725. return self.mImgList[:5]
  726. elif currIndex > len(self.mImgList)-3:
  727. return self.mImgList[-5:]
  728. else:
  729. return self.mImgList[currIndex - 2 : currIndex + 3]
  730. # Tzutalin 20160906 : Add file list and dock to move faster
  731. def fileitemDoubleClicked(self, item=None):
  732. self.currIndex = self.mImgList.index(ustr(os.path.join(os.path.abspath(self.dirname), item.text())))
  733. filename = self.mImgList[self.currIndex]
  734. if filename:
  735. self.mImgList5 = self.indexTo5Files(self.currIndex)
  736. # self.additems5(None)
  737. self.loadFile(filename)
  738. def iconitemDoubleClicked(self, item=None):
  739. self.currIndex = self.mImgList.index(ustr(os.path.join(item.toolTip())))
  740. filename = self.mImgList[self.currIndex]
  741. if filename:
  742. self.mImgList5 = self.indexTo5Files(self.currIndex)
  743. # self.additems5(None)
  744. self.loadFile(filename)
  745. def CanvasSizeChange(self):
  746. if len(self.mImgList) > 0:
  747. self.zoomWidget.setValue(self.zoomWidgetValue + self.imgsplider.value())
  748. # React to canvas signals.
  749. def shapeSelectionChanged(self, selected=False):
  750. if self._noSelectionSlot:
  751. self._noSelectionSlot = False
  752. else:
  753. shape = self.canvas.selectedShape
  754. if shape:
  755. self.shapesToItems[shape].setSelected(True)
  756. self.shapesToItemsbox[shape].setSelected(True) # ADD
  757. else:
  758. self.labelList.clearSelection()
  759. self.actions.delete.setEnabled(selected)
  760. self.actions.copy.setEnabled(selected)
  761. self.actions.edit.setEnabled(selected)
  762. self.actions.shapeLineColor.setEnabled(selected)
  763. self.actions.shapeFillColor.setEnabled(selected)
  764. self.actions.singleRere.setEnabled(selected)
  765. def addLabel(self, shape):
  766. shape.paintLabel = self.displayLabelOption.isChecked()
  767. item = HashableQListWidgetItem(shape.label)
  768. item.setFlags(item.flags() | Qt.ItemIsUserCheckable)
  769. item.setCheckState(Qt.Unchecked) if shape.difficult else item.setCheckState(Qt.Checked)
  770. # Checked means difficult is False
  771. # item.setBackground(generateColorByText(shape.label))
  772. self.itemsToShapes[item] = shape
  773. self.shapesToItems[shape] = item
  774. self.labelList.addItem(item)
  775. # print('item in add label is ',[(p.x(), p.y()) for p in shape.points], shape.label)
  776. # ADD for box
  777. item = HashableQListWidgetItem(str([(int(p.x()), int(p.y())) for p in shape.points]))
  778. self.itemsToShapesbox[item] = shape
  779. self.shapesToItemsbox[shape] = item
  780. self.BoxList.addItem(item)
  781. for action in self.actions.onShapesPresent:
  782. action.setEnabled(True)
  783. self.updateComboBox()
  784. def remLabel(self, shape):
  785. if shape is None:
  786. # print('rm empty label')
  787. return
  788. item = self.shapesToItems[shape]
  789. self.labelList.takeItem(self.labelList.row(item))
  790. del self.shapesToItems[shape]
  791. del self.itemsToShapes[item]
  792. self.updateComboBox()
  793. # ADD:
  794. item = self.shapesToItemsbox[shape]
  795. self.BoxList.takeItem(self.BoxList.row(item))
  796. del self.shapesToItemsbox[shape]
  797. del self.itemsToShapesbox[item]
  798. self.updateComboBox()
  799. def loadLabels(self, shapes):
  800. s = []
  801. for label, points, line_color, fill_color, difficult in shapes:
  802. shape = Shape(label=label)
  803. for x, y in points:
  804. # Ensure the labels are within the bounds of the image. If not, fix them.
  805. x, y, snapped = self.canvas.snapPointToCanvas(x, y)
  806. if snapped:
  807. self.setDirty()
  808. shape.addPoint(QPointF(x, y))
  809. shape.difficult = difficult
  810. shape.close()
  811. s.append(shape)
  812. # if line_color:
  813. # shape.line_color = QColor(*line_color)
  814. # else:
  815. # shape.line_color = generateColorByText(label)
  816. #
  817. # if fill_color:
  818. # shape.fill_color = QColor(*fill_color)
  819. # else:
  820. # shape.fill_color = generateColorByText(label)
  821. self.addLabel(shape)
  822. self.updateComboBox()
  823. self.canvas.loadShapes(s)
  824. def singleLabel(self, shape):
  825. if shape is None:
  826. # print('rm empty label')
  827. return
  828. item = self.shapesToItems[shape]
  829. item.setText(shape.label)
  830. self.updateComboBox()
  831. # ADD:
  832. item = self.shapesToItemsbox[shape]
  833. item.setText(str([(int(p.x()), int(p.y())) for p in shape.points]))
  834. self.updateComboBox()
  835. def updateComboBox(self):
  836. # Get the unique labels and add them to the Combobox.
  837. itemsTextList = [str(self.labelList.item(i).text()) for i in range(self.labelList.count())]
  838. uniqueTextList = list(set(itemsTextList))
  839. # Add a null row for showing all the labels
  840. uniqueTextList.append("")
  841. uniqueTextList.sort()
  842. # self.comboBox.update_items(uniqueTextList)
  843. def saveLabels(self, annotationFilePath, mode='Auto'):
  844. # Mode is Auto means that labels will be loaded from self.result_dic totally, which is the output of ocr model
  845. annotationFilePath = ustr(annotationFilePath)
  846. def format_shape(s):
  847. # print('s in saveLabels is ',s)
  848. return dict(label=s.label, # str
  849. line_color=s.line_color.getRgb(),
  850. fill_color=s.fill_color.getRgb(),
  851. points=[(p.x(), p.y()) for p in s.points], # QPonitF
  852. # add chris
  853. difficult=s.difficult) # bool
  854. shapes = [] if mode == 'Auto' else \
  855. [format_shape(shape) for shape in self.canvas.shapes]
  856. # Can add differrent annotation formats here
  857. for box in self.result_dic:
  858. trans_dic = {"label": box[1][0], "points": box[0], 'difficult': False}
  859. if trans_dic["label"] == "" and mode == 'Auto':
  860. continue
  861. shapes.append(trans_dic)
  862. try:
  863. trans_dic = []
  864. for box in shapes:
  865. trans_dic.append({"transcription": box['label'], "points": box['points'], 'difficult': box['difficult']})
  866. self.PPlabel[annotationFilePath] = trans_dic
  867. if mode == 'Auto':
  868. self.Cachelabel[annotationFilePath] = trans_dic
  869. # else:
  870. # self.labelFile.save(annotationFilePath, shapes, self.filePath, self.imageData,
  871. # self.lineColor.getRgb(), self.fillColor.getRgb())
  872. # print('Image:{0} -> Annotation:{1}'.format(self.filePath, annotationFilePath))
  873. return True
  874. except:
  875. self.errorMessage(u'Error saving label data')
  876. return False
  877. def copySelectedShape(self):
  878. self.addLabel(self.canvas.copySelectedShape())
  879. # fix copy and delete
  880. self.shapeSelectionChanged(True)
  881. def labelSelectionChanged(self):
  882. item = self.currentItem()
  883. self.labelList.scrollToItem(item, QAbstractItemView.EnsureVisible)
  884. if item and self.canvas.editing():
  885. self._noSelectionSlot = True
  886. self.canvas.selectShape(self.itemsToShapes[item])
  887. shape = self.itemsToShapes[item]
  888. def boxSelectionChanged(self):
  889. item = self.currentBox()
  890. self.BoxList.scrollToItem(item, QAbstractItemView.EnsureVisible)
  891. if item and self.canvas.editing():
  892. self._noSelectionSlot = True
  893. self.canvas.selectShape(self.itemsToShapesbox[item])
  894. shape = self.itemsToShapesbox[item]
  895. def labelItemChanged(self, item):
  896. shape = self.itemsToShapes[item]
  897. label = item.text()
  898. if label != shape.label:
  899. shape.label = item.text()
  900. # shape.line_color = generateColorByText(shape.label)
  901. self.setDirty()
  902. elif not ((item.checkState()== Qt.Unchecked) ^ (not shape.difficult)):
  903. shape.difficult = True if item.checkState() == Qt.Unchecked else False
  904. self.setDirty()
  905. else: # User probably changed item visibility
  906. self.canvas.setShapeVisible(shape, True) # item.checkState() == Qt.Checked
  907. # self.actions.save.setEnabled(True)
  908. # Callback functions:
  909. def newShape(self, value=True):
  910. """Pop-up and give focus to the label editor.
  911. position MUST be in global coordinates.
  912. """
  913. if len(self.labelHist) > 0:
  914. self.labelDialog = LabelDialog(
  915. parent=self, listItem=self.labelHist)
  916. if value:
  917. text = self.labelDialog.popUp(text=self.prevLabelText)
  918. self.lastLabel = text
  919. else:
  920. text = self.prevLabelText
  921. if text is not None:
  922. self.prevLabelText = self.stringBundle.getString('tempLabel')
  923. # generate_color = generateColorByText(text)
  924. shape = self.canvas.setLastLabel(text, None, None)#generate_color, generate_color
  925. self.addLabel(shape)
  926. if self.beginner(): # Switch to edit mode.
  927. self.canvas.setEditing(True)
  928. self.actions.create.setEnabled(True)
  929. else:
  930. self.actions.editMode.setEnabled(True)
  931. self.setDirty()
  932. else:
  933. # self.canvas.undoLastLine()
  934. self.canvas.resetAllLines()
  935. def scrollRequest(self, delta, orientation):
  936. units = - delta / (8 * 15)
  937. bar = self.scrollBars[orientation]
  938. bar.setValue(bar.value() + bar.singleStep() * units)
  939. def setZoom(self, value):
  940. self.actions.fitWidth.setChecked(False)
  941. self.actions.fitWindow.setChecked(False)
  942. self.zoomMode = self.MANUAL_ZOOM
  943. self.zoomWidget.setValue(value)
  944. def addZoom(self, increment=10):
  945. self.setZoom(self.zoomWidget.value() + increment)
  946. def zoomRequest(self, delta):
  947. # get the current scrollbar positions
  948. # calculate the percentages ~ coordinates
  949. h_bar = self.scrollBars[Qt.Horizontal]
  950. v_bar = self.scrollBars[Qt.Vertical]
  951. # get the current maximum, to know the difference after zooming
  952. h_bar_max = h_bar.maximum()
  953. v_bar_max = v_bar.maximum()
  954. # get the cursor position and canvas size
  955. # calculate the desired movement from 0 to 1
  956. # where 0 = move left
  957. # 1 = move right
  958. # up and down analogous
  959. cursor = QCursor()
  960. pos = cursor.pos()
  961. relative_pos = QWidget.mapFromGlobal(self, pos)
  962. cursor_x = relative_pos.x()
  963. cursor_y = relative_pos.y()
  964. w = self.scrollArea.width()
  965. h = self.scrollArea.height()
  966. # the scaling from 0 to 1 has some padding
  967. # you don't have to hit the very leftmost pixel for a maximum-left movement
  968. margin = 0.1
  969. move_x = (cursor_x - margin * w) / (w - 2 * margin * w)
  970. move_y = (cursor_y - margin * h) / (h - 2 * margin * h)
  971. # clamp the values from 0 to 1
  972. move_x = min(max(move_x, 0), 1)
  973. move_y = min(max(move_y, 0), 1)
  974. # zoom in
  975. units = delta / (8 * 15)
  976. scale = 10
  977. self.addZoom(scale * units)
  978. # get the difference in scrollbar values
  979. # this is how far we can move
  980. d_h_bar_max = h_bar.maximum() - h_bar_max
  981. d_v_bar_max = v_bar.maximum() - v_bar_max
  982. # get the new scrollbar values
  983. new_h_bar_value = h_bar.value() + move_x * d_h_bar_max
  984. new_v_bar_value = v_bar.value() + move_y * d_v_bar_max
  985. h_bar.setValue(new_h_bar_value)
  986. v_bar.setValue(new_v_bar_value)
  987. def setFitWindow(self, value=True):
  988. if value:
  989. self.actions.fitWidth.setChecked(False)
  990. self.zoomMode = self.FIT_WINDOW if value else self.MANUAL_ZOOM
  991. self.adjustScale()
  992. def setFitWidth(self, value=True):
  993. if value:
  994. self.actions.fitWindow.setChecked(False)
  995. self.zoomMode = self.FIT_WIDTH if value else self.MANUAL_ZOOM
  996. self.adjustScale()
  997. def togglePolygons(self, value):
  998. for item, shape in self.itemsToShapes.items():
  999. self.canvas.setShapeVisible(shape, value)
  1000. def loadFile(self, filePath=None):
  1001. """Load the specified file, or the last opened file if None."""
  1002. self.resetState()
  1003. self.canvas.setEnabled(False)
  1004. if filePath is None:
  1005. filePath = self.settings.get(SETTING_FILENAME)
  1006. # Make sure that filePath is a regular python string, rather than QString
  1007. filePath = ustr(filePath)
  1008. # Fix bug: An index error after select a directory when open a new file.
  1009. unicodeFilePath = ustr(filePath)
  1010. # unicodeFilePath = os.path.abspath(unicodeFilePath)
  1011. # Tzutalin 20160906 : Add file list and dock to move faster
  1012. # Highlight the file item
  1013. if unicodeFilePath and self.fileListWidget.count() > 0:
  1014. if unicodeFilePath in self.mImgList:
  1015. index = self.mImgList.index(unicodeFilePath)
  1016. fileWidgetItem = self.fileListWidget.item(index)
  1017. print('unicodeFilePath is', unicodeFilePath)
  1018. fileWidgetItem.setSelected(True)
  1019. ###
  1020. self.iconlist.clear()
  1021. self.additems5(None)
  1022. for i in range(5):
  1023. item_tooltip = self.iconlist.item(i).toolTip()
  1024. # print(i,"---",item_tooltip)
  1025. if item_tooltip == ustr(filePath):
  1026. titem = self.iconlist.item(i)
  1027. titem.setSelected(True)
  1028. self.iconlist.scrollToItem(titem)
  1029. break
  1030. else:
  1031. self.fileListWidget.clear()
  1032. self.mImgList.clear()
  1033. self.iconlist.clear()
  1034. # if unicodeFilePath and self.iconList.count() > 0:
  1035. # if unicodeFilePath in self.mImgList:
  1036. if unicodeFilePath and os.path.exists(unicodeFilePath):
  1037. self.canvas.verified = False
  1038. cvimg = cv2.imdecode(np.fromfile(unicodeFilePath, dtype=np.uint8), 1)
  1039. height, width, depth = cvimg.shape
  1040. cvimg = cv2.cvtColor(cvimg, cv2.COLOR_BGR2RGB)
  1041. image = QImage(cvimg.data, width, height, width * depth, QImage.Format_RGB888)
  1042. if image.isNull():
  1043. self.errorMessage(u'Error opening file',
  1044. u"<p>Make sure <i>%s</i> is a valid image file." % unicodeFilePath)
  1045. self.status("Error reading %s" % unicodeFilePath)
  1046. return False
  1047. self.status("Loaded %s" % os.path.basename(unicodeFilePath))
  1048. self.image = image
  1049. self.filePath = unicodeFilePath
  1050. self.canvas.loadPixmap(QPixmap.fromImage(image))
  1051. if self.validFilestate(filePath) is True:
  1052. self.setClean()
  1053. else:
  1054. self.dirty = False
  1055. self.actions.save.setEnabled(True)
  1056. self.canvas.setEnabled(True)
  1057. self.adjustScale(initial=True)
  1058. self.paintCanvas()
  1059. self.addRecentFile(self.filePath)
  1060. self.toggleActions(True)
  1061. self.showBoundingBoxFromPPlabel(filePath)
  1062. self.setWindowTitle(__appname__ + ' ' + filePath)
  1063. # Default : select last item if there is at least one item
  1064. if self.labelList.count():
  1065. self.labelList.setCurrentItem(self.labelList.item(self.labelList.count() - 1))
  1066. self.labelList.item(self.labelList.count() - 1).setSelected(True)
  1067. self.canvas.setFocus(True)
  1068. return True
  1069. return False
  1070. def showBoundingBoxFromPPlabel(self, filePath):
  1071. imgidx = self.getImglabelidx(filePath)
  1072. if imgidx not in self.PPlabel.keys():
  1073. return
  1074. shapes = []
  1075. for box in self.PPlabel[imgidx]:
  1076. shapes.append((box['transcription'], box['points'], None, None, box['difficult']))
  1077. print(shapes)
  1078. self.loadLabels(shapes)
  1079. self.canvas.verified = False
  1080. def validFilestate(self, filePath):
  1081. if filePath not in self.fileStatedict.keys():
  1082. return None
  1083. elif self.fileStatedict[filePath] == 1:
  1084. return True
  1085. else:
  1086. return False
  1087. def resizeEvent(self, event):
  1088. if self.canvas and not self.image.isNull() \
  1089. and self.zoomMode != self.MANUAL_ZOOM:
  1090. self.adjustScale()
  1091. super(MainWindow, self).resizeEvent(event)
  1092. def paintCanvas(self):
  1093. assert not self.image.isNull(), "cannot paint null image"
  1094. self.canvas.scale = 0.01 * self.zoomWidget.value()
  1095. self.canvas.adjustSize()
  1096. self.canvas.update()
  1097. def adjustScale(self, initial=False):
  1098. value = self.scalers[self.FIT_WINDOW if initial else self.zoomMode]()
  1099. self.zoomWidget.setValue(int(100 * value))
  1100. def scaleFitWindow(self):
  1101. """Figure out the size of the pixmap in order to fit the main widget."""
  1102. e = 2.0 # So that no scrollbars are generated.
  1103. w1 = self.centralWidget().width() - e
  1104. h1 = self.centralWidget().height() - e -110
  1105. a1 = w1 / h1
  1106. # Calculate a new scale value based on the pixmap's aspect ratio.
  1107. w2 = self.canvas.pixmap.width() - 0.0
  1108. h2 = self.canvas.pixmap.height() - 0.0
  1109. a2 = w2 / h2
  1110. return w1 / w2 if a2 >= a1 else h1 / h2
  1111. def scaleFitWidth(self):
  1112. # The epsilon does not seem to work too well here.
  1113. w = self.centralWidget().width() - 2.0
  1114. return w / self.canvas.pixmap.width()
  1115. def closeEvent(self, event):
  1116. if not self.mayContinue():
  1117. event.ignore()
  1118. else:
  1119. settings = self.settings
  1120. # If it loads images from dir, don't load it at the begining
  1121. if self.dirname is None:
  1122. settings[SETTING_FILENAME] = self.filePath if self.filePath else ''
  1123. else:
  1124. settings[SETTING_FILENAME] = ''
  1125. settings[SETTING_WIN_SIZE] = self.size()
  1126. settings[SETTING_WIN_POSE] = self.pos()
  1127. settings[SETTING_WIN_STATE] = self.saveState()
  1128. settings[SETTING_LINE_COLOR] = self.lineColor
  1129. settings[SETTING_FILL_COLOR] = self.fillColor
  1130. settings[SETTING_RECENT_FILES] = self.recentFiles
  1131. settings[SETTING_ADVANCE_MODE] = not self._beginner
  1132. if self.defaultSaveDir and os.path.exists(self.defaultSaveDir):
  1133. settings[SETTING_SAVE_DIR] = ustr(self.defaultSaveDir)
  1134. else:
  1135. settings[SETTING_SAVE_DIR] = ''
  1136. if self.lastOpenDir and os.path.exists(self.lastOpenDir):
  1137. settings[SETTING_LAST_OPEN_DIR] = self.lastOpenDir
  1138. else:
  1139. settings[SETTING_LAST_OPEN_DIR] = ''
  1140. settings[SETTING_PAINT_LABEL] = self.displayLabelOption.isChecked()
  1141. settings[SETTING_DRAW_SQUARE] = self.drawSquaresOption.isChecked()
  1142. settings.save()
  1143. try:
  1144. self.saveLabelFile()
  1145. except:
  1146. pass
  1147. def loadRecent(self, filename):
  1148. if self.mayContinue():
  1149. self.loadFile(filename)
  1150. def scanAllImages(self, folderPath):
  1151. extensions = ['.%s' % fmt.data().decode("ascii").lower() for fmt in QImageReader.supportedImageFormats()]
  1152. images = []
  1153. for file in os.listdir(folderPath):
  1154. if file.lower().endswith(tuple(extensions)):
  1155. relativePath = os.path.join(folderPath, file)
  1156. path = ustr(os.path.abspath(relativePath))
  1157. images.append(path)
  1158. natural_sort(images, key=lambda x: x.lower())
  1159. return images
  1160. def openDirDialog(self, _value=False, dirpath=None, silent=False):
  1161. if not self.mayContinue():
  1162. return
  1163. defaultOpenDirPath = dirpath if dirpath else '.'
  1164. if self.lastOpenDir and os.path.exists(self.lastOpenDir):
  1165. defaultOpenDirPath = self.lastOpenDir
  1166. else:
  1167. defaultOpenDirPath = os.path.dirname(self.filePath) if self.filePath else '.'
  1168. if silent != True:
  1169. targetDirPath = ustr(QFileDialog.getExistingDirectory(self,
  1170. '%s - Open Directory' % __appname__,
  1171. defaultOpenDirPath,
  1172. QFileDialog.ShowDirsOnly | QFileDialog.DontResolveSymlinks))
  1173. else:
  1174. targetDirPath = ustr(defaultOpenDirPath)
  1175. self.lastOpenDir = targetDirPath
  1176. self.importDirImages(targetDirPath)
  1177. def importDirImages(self, dirpath, isDelete = False):
  1178. if not self.mayContinue() or not dirpath:
  1179. return
  1180. if self.defaultSaveDir and self.defaultSaveDir != dirpath:
  1181. self.saveLabelFile()
  1182. if not isDelete:
  1183. self.loadFilestate(dirpath)
  1184. self.PPlabelpath = dirpath+ '/Label.txt'
  1185. self.PPlabel = self.loadLabelFile(self.PPlabelpath)
  1186. self.Cachelabelpath = dirpath + '/Cache.cach'
  1187. self.Cachelabel = self.loadLabelFile(self.Cachelabelpath)
  1188. if self.Cachelabel:
  1189. self.PPlabel = dict(self.Cachelabel, **self.PPlabel)
  1190. self.lastOpenDir = dirpath
  1191. self.dirname = dirpath
  1192. self.defaultSaveDir = dirpath
  1193. self.statusBar().showMessage('%s started. Annotation will be saved to %s' %
  1194. (__appname__, self.defaultSaveDir))
  1195. self.statusBar().show()
  1196. self.filePath = None
  1197. self.fileListWidget.clear()
  1198. self.mImgList = self.scanAllImages(dirpath)
  1199. self.mImgList5 = self.mImgList[:5]
  1200. self.openNextImg()
  1201. doneicon = newIcon('done')
  1202. closeicon = newIcon('close')
  1203. for imgPath in self.mImgList:
  1204. filename = os.path.basename(imgPath)
  1205. if self.validFilestate(imgPath) is True:
  1206. item = QListWidgetItem(doneicon, filename)
  1207. else:
  1208. item = QListWidgetItem(closeicon, filename)
  1209. self.fileListWidget.addItem(item)
  1210. print('DirPath in importDirImages is', dirpath)
  1211. self.iconlist.clear()
  1212. self.additems5(dirpath)
  1213. self.changeFileFolder = True
  1214. self.haveAutoReced = False
  1215. self.AutoRecognition.setEnabled(True)
  1216. self.reRecogButton.setEnabled(True)
  1217. self.actions.AutoRec.setEnabled(True)
  1218. self.actions.reRec.setEnabled(True)
  1219. def openPrevImg(self, _value=False):
  1220. if len(self.mImgList) <= 0:
  1221. return
  1222. if self.filePath is None:
  1223. return
  1224. currIndex = self.mImgList.index(self.filePath)
  1225. self.mImgList5 = self.mImgList[:5]
  1226. if currIndex - 1 >= 0:
  1227. filename = self.mImgList[currIndex - 1]
  1228. self.mImgList5 = self.indexTo5Files(currIndex - 1)
  1229. if filename:
  1230. self.loadFile(filename)
  1231. def openNextImg(self, _value=False):
  1232. if not self.mayContinue():
  1233. return
  1234. if len(self.mImgList) <= 0:
  1235. return
  1236. filename = None
  1237. if self.filePath is None:
  1238. filename = self.mImgList[0]
  1239. self.mImgList5 = self.mImgList[:5]
  1240. else:
  1241. currIndex = self.mImgList.index(self.filePath)
  1242. if currIndex + 1 < len(self.mImgList):
  1243. filename = self.mImgList[currIndex + 1]
  1244. self.mImgList5 = self.indexTo5Files(currIndex + 1)
  1245. else:
  1246. self.mImgList5 = self.indexTo5Files(currIndex)
  1247. if filename:
  1248. print('file name in openNext is ',filename)
  1249. self.loadFile(filename)
  1250. def updateFileListIcon(self, filename):
  1251. pass
  1252. def saveFile(self, _value=False, mode='Manual'):
  1253. # Manual mode is used for users click "Save" manually,which will change the state of the image
  1254. if self.filePath:
  1255. imgidx = self.getImglabelidx(self.filePath)
  1256. self._saveFile(imgidx, mode=mode)
  1257. def saveFileAs(self, _value=False):
  1258. assert not self.image.isNull(), "cannot save empty image"
  1259. self._saveFile(self.saveFileDialog())
  1260. def saveFileDialog(self, removeExt=True):
  1261. caption = '%s - Choose File' % __appname__
  1262. filters = 'File (*%s)' % LabelFile.suffix
  1263. openDialogPath = self.currentPath()
  1264. dlg = QFileDialog(self, caption, openDialogPath, filters)
  1265. dlg.setDefaultSuffix(LabelFile.suffix[1:])
  1266. dlg.setAcceptMode(QFileDialog.AcceptSave)
  1267. filenameWithoutExtension = os.path.splitext(self.filePath)[0]
  1268. dlg.selectFile(filenameWithoutExtension)
  1269. dlg.setOption(QFileDialog.DontUseNativeDialog, False)
  1270. if dlg.exec_():
  1271. fullFilePath = ustr(dlg.selectedFiles()[0])
  1272. if removeExt:
  1273. return os.path.splitext(fullFilePath)[0] # Return file path without the extension.
  1274. else:
  1275. return fullFilePath
  1276. return ''
  1277. def _saveFile(self, annotationFilePath, mode='Manual'):
  1278. if mode == 'Manual':
  1279. if annotationFilePath and self.saveLabels(annotationFilePath, mode=mode):
  1280. self.setClean()
  1281. self.statusBar().showMessage('Saved to %s' % annotationFilePath)
  1282. self.statusBar().show()
  1283. currIndex = self.mImgList.index(self.filePath)
  1284. item = self.fileListWidget.item(currIndex)
  1285. item.setIcon(newIcon('done'))
  1286. self.fileStatedict[self.filePath] = 1
  1287. if len(self.fileStatedict)%self.autoSaveNum ==0:
  1288. self.saveFilestate()
  1289. self.savePPlabel(mode='Auto')
  1290. self.fileListWidget.insertItem(int(currIndex), item)
  1291. self.openNextImg()
  1292. self.actions.saveRec.setEnabled(True)
  1293. elif mode == 'Auto':
  1294. if annotationFilePath and self.saveLabels(annotationFilePath, mode=mode):
  1295. self.setClean()
  1296. self.statusBar().showMessage('Saved to %s' % annotationFilePath)
  1297. self.statusBar().show()
  1298. def closeFile(self, _value=False):
  1299. if not self.mayContinue():
  1300. return
  1301. self.resetState()
  1302. self.setClean()
  1303. self.toggleActions(False)
  1304. self.canvas.setEnabled(False)
  1305. self.actions.saveAs.setEnabled(False)
  1306. def deleteImg(self):
  1307. deletePath = self.filePath
  1308. if deletePath is not None:
  1309. deleteInfo = self.deleteImgDialog()
  1310. if deleteInfo == QMessageBox.Yes:
  1311. if platform.system() == 'Windows':
  1312. from win32com.shell import shell, shellcon
  1313. shell.SHFileOperation((0, shellcon.FO_DELETE, deletePath, None,
  1314. shellcon.FOF_SILENT | shellcon.FOF_ALLOWUNDO | shellcon.FOF_NOCONFIRMATION,
  1315. None, None))
  1316. # linux
  1317. elif platform.system() == 'Linux':
  1318. cmd = 'trash ' + deletePath
  1319. os.system(cmd)
  1320. # macOS
  1321. elif platform.system() == 'Darwin':
  1322. import subprocess
  1323. absPath = os.path.abspath(deletePath).replace('\\', '\\\\').replace('"', '\\"')
  1324. cmd = ['osascript', '-e',
  1325. 'tell app "Finder" to move {the POSIX file "' + absPath + '"} to trash']
  1326. print(cmd)
  1327. subprocess.call(cmd, stdout=open(os.devnull, 'w'))
  1328. if self.filePath in self.fileStatedict.keys():
  1329. self.fileStatedict.pop(self.filePath)
  1330. imgidx = self.getImglabelidx(self.filePath)
  1331. if imgidx in self.PPlabel.keys():
  1332. self.PPlabel.pop(imgidx)
  1333. self.openNextImg()
  1334. self.importDirImages(self.lastOpenDir, isDelete=True)
  1335. def deleteImgDialog(self):
  1336. yes, cancel = QMessageBox.Yes, QMessageBox.Cancel
  1337. msg = u'The image will be deleted to the recycle bin'
  1338. return QMessageBox.warning(self, u'Attention', msg, yes | cancel)
  1339. def resetAll(self):
  1340. self.settings.reset()
  1341. self.close()
  1342. proc = QProcess()
  1343. proc.startDetached(os.path.abspath(__file__))
  1344. def mayContinue(self): #
  1345. if not self.dirty:
  1346. return True
  1347. else:
  1348. discardChanges = self.discardChangesDialog()
  1349. if discardChanges == QMessageBox.No:
  1350. return True
  1351. elif discardChanges == QMessageBox.Yes:
  1352. self.saveFile()
  1353. return True
  1354. else:
  1355. return False
  1356. def discardChangesDialog(self):
  1357. yes, no, cancel = QMessageBox.Yes, QMessageBox.No, QMessageBox.Cancel
  1358. msg = u'You have unsaved changes, would you like to save them and proceed?\nClick "No" to undo all changes.'
  1359. return QMessageBox.warning(self, u'Attention', msg, yes | no | cancel)
  1360. def errorMessage(self, title, message):
  1361. return QMessageBox.critical(self, title,
  1362. '<p><b>%s</b></p>%s' % (title, message))
  1363. def currentPath(self):
  1364. return os.path.dirname(self.filePath) if self.filePath else '.'
  1365. def chooseColor1(self):
  1366. color = self.colorDialog.getColor(self.lineColor, u'Choose line color',
  1367. default=DEFAULT_LINE_COLOR)
  1368. if color:
  1369. self.lineColor = color
  1370. Shape.line_color = color
  1371. self.canvas.setDrawingColor(color)
  1372. self.canvas.update()
  1373. self.setDirty()
  1374. def deleteSelectedShape(self):
  1375. self.remLabel(self.canvas.deleteSelected())
  1376. self.setDirty()
  1377. if self.noShapes():
  1378. for action in self.actions.onShapesPresent:
  1379. action.setEnabled(False)
  1380. def chshapeLineColor(self):
  1381. color = self.colorDialog.getColor(self.lineColor, u'Choose line color',
  1382. default=DEFAULT_LINE_COLOR)
  1383. if color:
  1384. self.canvas.selectedShape.line_color = color
  1385. self.canvas.update()
  1386. self.setDirty()
  1387. def chshapeFillColor(self):
  1388. color = self.colorDialog.getColor(self.fillColor, u'Choose fill color',
  1389. default=DEFAULT_FILL_COLOR)
  1390. if color:
  1391. self.canvas.selectedShape.fill_color = color
  1392. self.canvas.update()
  1393. self.setDirty()
  1394. def copyShape(self):
  1395. self.canvas.endMove(copy=True)
  1396. self.addLabel(self.canvas.selectedShape)
  1397. self.setDirty()
  1398. def moveShape(self):
  1399. self.canvas.endMove(copy=False)
  1400. self.setDirty()
  1401. def loadPredefinedClasses(self, predefClassesFile):
  1402. if os.path.exists(predefClassesFile) is True:
  1403. with codecs.open(predefClassesFile, 'r', 'utf8') as f:
  1404. for line in f:
  1405. line = line.strip()
  1406. if self.labelHist is None:
  1407. self.labelHist = [line]
  1408. else:
  1409. self.labelHist.append(line)
  1410. def togglePaintLabelsOption(self):
  1411. for shape in self.canvas.shapes:
  1412. shape.paintLabel = self.displayLabelOption.isChecked()
  1413. def toogleDrawSquare(self):
  1414. self.canvas.setDrawingShapeToSquare(self.drawSquaresOption.isChecked())
  1415. def additems(self, dirpath):
  1416. for file in self.mImgList:
  1417. pix = QPixmap(file)
  1418. _, filename = os.path.split(file)
  1419. filename, _ = os.path.splitext(filename)
  1420. item = QListWidgetItem(QIcon(pix.scaled(100, 100, Qt.IgnoreAspectRatio, Qt.FastTransformation)),
  1421. filename[:10])
  1422. item.setToolTip(file)
  1423. self.iconlist.addItem(item)
  1424. def additems5(self, dirpath):
  1425. for file in self.mImgList5:
  1426. pix = QPixmap(file)
  1427. _, filename = os.path.split(file)
  1428. filename, _ = os.path.splitext(filename)
  1429. pfilename = filename[:10]
  1430. if len(pfilename) < 10:
  1431. lentoken = 12 - len(pfilename)
  1432. prelen = lentoken // 2
  1433. bfilename = prelen * " " + pfilename + (lentoken - prelen) * " "
  1434. # item = QListWidgetItem(QIcon(pix.scaled(100, 100, Qt.KeepAspectRatio, Qt.SmoothTransformation)),filename[:10])
  1435. item = QListWidgetItem(QIcon(pix.scaled(100, 100, Qt.IgnoreAspectRatio, Qt.FastTransformation)),pfilename)
  1436. # item.setForeground(QBrush(Qt.white))
  1437. item.setToolTip(file)
  1438. self.iconlist.addItem(item)
  1439. owidth = 0
  1440. for index in range(len(self.mImgList5)):
  1441. item = self.iconlist.item(index)
  1442. itemwidget = self.iconlist.visualItemRect(item)
  1443. owidth += itemwidget.width()
  1444. self.iconlist.setMinimumWidth(owidth + 50)
  1445. def getImglabelidx(self, filePath):
  1446. if platform.system()=='Windows':
  1447. spliter = '\\'
  1448. else:
  1449. spliter = '/'
  1450. filepathsplit = filePath.split(spliter)[-2:]
  1451. return filepathsplit[0] + '/' + filepathsplit[1]
  1452. def autoRecognition(self):
  1453. assert self.mImgList is not None
  1454. print('Using model from ', self.model)
  1455. uncheckedList = [i for i in self.mImgList if i not in self.fileStatedict.keys()]
  1456. self.autoDialog = AutoDialog(parent=self, ocr=self.ocr, mImgList=uncheckedList, lenbar=len(uncheckedList))
  1457. self.autoDialog.popUp()
  1458. self.currIndex=len(self.mImgList)
  1459. self.loadFile(self.filePath) # ADD
  1460. self.haveAutoReced = True
  1461. self.AutoRecognition.setEnabled(False)
  1462. self.actions.AutoRec.setEnabled(False)
  1463. self.setDirty()
  1464. self.saveCacheLabel()
  1465. def reRecognition(self):
  1466. img = cv2.imread(self.filePath)
  1467. # org_box = [dic['points'] for dic in self.PPlabel[self.getImglabelidx(self.filePath)]]
  1468. if self.canvas.shapes:
  1469. self.result_dic = []
  1470. rec_flag = 0
  1471. for shape in self.canvas.shapes:
  1472. box = [[int(p.x()), int(p.y())] for p in shape.points]
  1473. assert len(box) == 4
  1474. img_crop = get_rotate_crop_image(img, np.array(box, np.float32))
  1475. if img_crop is None:
  1476. msg = 'Can not recognise the detection box in ' + self.filePath + '. Please change manually'
  1477. QMessageBox.information(self, "Information", msg)
  1478. return
  1479. result = self.ocr.ocr(img_crop, cls=True, det=False)
  1480. if result[0][0] != '':
  1481. result.insert(0, box)
  1482. print('result in reRec is ', result)
  1483. self.result_dic.append(result)
  1484. if result[1][0] == shape.label:
  1485. print('label no change')
  1486. else:
  1487. rec_flag += 1
  1488. if len(self.result_dic) > 0 and rec_flag > 0:
  1489. self.saveFile(mode='Auto')
  1490. self.loadFile(self.filePath)
  1491. self.setDirty()
  1492. elif len(self.result_dic) == len(self.canvas.shapes) and rec_flag == 0:
  1493. QMessageBox.information(self, "Information", "The recognition result remains unchanged!")
  1494. else:
  1495. print('Can not recgonise in ', self.filePath)
  1496. else:
  1497. QMessageBox.information(self, "Information", "Draw a box!")
  1498. def singleRerecognition(self):
  1499. img = cv2.imread(self.filePath)
  1500. shape = self.canvas.selectedShape
  1501. box = [[int(p.x()), int(p.y())] for p in shape.points]
  1502. assert len(box) == 4
  1503. img_crop = get_rotate_crop_image(img, np.array(box, np.float32))
  1504. if img_crop is None:
  1505. msg = 'Can not recognise the detection box in ' + self.filePath + '. Please change manually'
  1506. QMessageBox.information(self, "Information", msg)
  1507. return
  1508. result = self.ocr.ocr(img_crop, cls=True, det=False)
  1509. if result[0][0] != '':
  1510. result.insert(0, box)
  1511. print('result in reRec is ', result)
  1512. if result[1][0] == shape.label:
  1513. print('label no change')
  1514. else:
  1515. shape.label = result[1][0]
  1516. self.singleLabel(shape)
  1517. self.setDirty()
  1518. print(box)
  1519. def autolcm(self):
  1520. vbox = QVBoxLayout()
  1521. hbox = QHBoxLayout()
  1522. self.panel = QLabel()
  1523. self.panel.setText(self.stringBundle.getString('choseModelLg'))
  1524. self.panel.setAlignment(Qt.AlignLeft)
  1525. self.comboBox = QComboBox()
  1526. self.comboBox.setObjectName("comboBox")
  1527. self.comboBox.addItems(['Chinese & English', 'English', 'French', 'German', 'Korean', 'Japanese'])
  1528. vbox.addWidget(self.panel)
  1529. vbox.addWidget(self.comboBox)
  1530. self.dialog = QDialog()
  1531. self.dialog.resize(300, 100)
  1532. self.okBtn = QPushButton(self.stringBundle.getString('ok'))
  1533. self.cancelBtn = QPushButton(self.stringBundle.getString('cancel'))
  1534. self.okBtn.clicked.connect(self.modelChoose)
  1535. self.cancelBtn.clicked.connect(self.cancel)
  1536. self.dialog.setWindowTitle(self.stringBundle.getString('choseModelLg'))
  1537. hbox.addWidget(self.okBtn)
  1538. hbox.addWidget(self.cancelBtn)
  1539. vbox.addWidget(self.panel)
  1540. vbox.addLayout(hbox)
  1541. self.dialog.setLayout(vbox)
  1542. self.dialog.setWindowModality(Qt.ApplicationModal)
  1543. self.dialog.exec_()
  1544. if self.filePath:
  1545. self.AutoRecognition.setEnabled(True)
  1546. self.actions.AutoRec.setEnabled(True)
  1547. def modelChoose(self):
  1548. print(self.comboBox.currentText())
  1549. lg_idx = {'Chinese & English': 'ch', 'English': 'en', 'French': 'french', 'German': 'german',
  1550. 'Korean': 'korean', 'Japanese': 'japan'}
  1551. del self.ocr
  1552. self.ocr = PaddleOCR(use_pdserving=False, use_angle_cls=True, det=True, cls=True, use_gpu=False,
  1553. lang=lg_idx[self.comboBox.currentText()])
  1554. self.dialog.close()
  1555. def cancel(self):
  1556. self.dialog.close()
  1557. def loadFilestate(self, saveDir):
  1558. self.fileStatepath = saveDir + '/fileState.txt'
  1559. self.fileStatedict = {}
  1560. if not os.path.exists(self.fileStatepath):
  1561. f = open(self.fileStatepath, 'w', encoding='utf-8')
  1562. else:
  1563. with open(self.fileStatepath, 'r', encoding='utf-8') as f:
  1564. states = f.readlines()
  1565. for each in states:
  1566. file, state = each.split('\t')
  1567. self.fileStatedict[file] = 1
  1568. self.actions.saveLabel.setEnabled(True)
  1569. self.actions.saveRec.setEnabled(True)
  1570. def saveFilestate(self):
  1571. with open(self.fileStatepath, 'w', encoding='utf-8') as f:
  1572. for key in self.fileStatedict:
  1573. f.write(key + '\t')
  1574. f.write(str(self.fileStatedict[key]) + '\n')
  1575. def loadLabelFile(self, labelpath):
  1576. labeldict = {}
  1577. if not os.path.exists(labelpath):
  1578. f = open(labelpath, 'w', encoding='utf-8')
  1579. else:
  1580. with open(labelpath, 'r', encoding='utf-8') as f:
  1581. data = f.readlines()
  1582. for each in data:
  1583. file, label = each.split('\t')
  1584. if label:
  1585. label = label.replace('false', 'False')
  1586. label = label.replace('true', 'True')
  1587. labeldict[file] = eval(label)
  1588. else:
  1589. labeldict[file] = []
  1590. return labeldict
  1591. def savePPlabel(self,mode='Manual'):
  1592. savedfile = [self.getImglabelidx(i) for i in self.fileStatedict.keys()]
  1593. with open(self.PPlabelpath, 'w', encoding='utf-8') as f:
  1594. for key in self.PPlabel:
  1595. if key in savedfile and self.PPlabel[key] != []:
  1596. f.write(key + '\t')
  1597. f.write(json.dumps(self.PPlabel[key], ensure_ascii=False) + '\n')
  1598. if mode=='Manual':
  1599. msg = 'Images that have been checked are saved in '+ self.PPlabelpath
  1600. QMessageBox.information(self, "Information", msg)
  1601. def saveCacheLabel(self):
  1602. with open(self.Cachelabelpath, 'w', encoding='utf-8') as f:
  1603. for key in self.Cachelabel:
  1604. f.write(key + '\t')
  1605. f.write(json.dumps(self.Cachelabel[key], ensure_ascii=False) + '\n')
  1606. def saveLabelFile(self):
  1607. self.saveFilestate()
  1608. self.savePPlabel()
  1609. def saveRecResult(self):
  1610. if None in [self.PPlabelpath, self.PPlabel, self.fileStatedict]:
  1611. QMessageBox.information(self, "Information", "Save file first")
  1612. return
  1613. rec_gt_dir = os.path.dirname(self.PPlabelpath) + '/rec_gt.txt'
  1614. crop_img_dir = os.path.dirname(self.PPlabelpath) + '/crop_img/'
  1615. ques_img = []
  1616. if not os.path.exists(crop_img_dir):
  1617. os.mkdir(crop_img_dir)
  1618. with open(rec_gt_dir, 'w', encoding='utf-8') as f:
  1619. for key in self.fileStatedict:
  1620. idx = self.getImglabelidx(key)
  1621. try:
  1622. img = cv2.imread(key)
  1623. for i, label in enumerate(self.PPlabel[idx]):
  1624. if label['difficult']: continue
  1625. img_crop = get_rotate_crop_image(img, np.array(label['points'], np.float32))
  1626. img_name = os.path.splitext(os.path.basename(idx))[0] + '_crop_'+str(i)+'.jpg'
  1627. cv2.imwrite(crop_img_dir+img_name, img_crop)
  1628. f.write('crop_img/'+ img_name + '\t')
  1629. f.write(label['transcription'] + '\n')
  1630. except Exception as e:
  1631. ques_img.append(key)
  1632. print("Can not read image ",e)
  1633. if ques_img:
  1634. QMessageBox.information(self, "Information", "The following images can not be saved, "
  1635. "please check the image path and labels.\n" + "".join(str(i)+'\n' for i in ques_img))
  1636. QMessageBox.information(self, "Information", "Cropped images have been saved in "+str(crop_img_dir))
  1637. def speedChoose(self):
  1638. if self.labelDialogOption.isChecked():
  1639. self.canvas.newShape.disconnect()
  1640. self.canvas.newShape.connect(partial(self.newShape, True))
  1641. else:
  1642. self.canvas.newShape.disconnect()
  1643. self.canvas.newShape.connect(partial(self.newShape, False))
  1644. def inverted(color):
  1645. return QColor(*[255 - v for v in color.getRgb()])
  1646. def read(filename, default=None):
  1647. try:
  1648. with open(filename, 'rb') as f:
  1649. return f.read()
  1650. except:
  1651. return default
  1652. def get_main_app(argv=[]):
  1653. """
  1654. Standard boilerplate Qt application code.
  1655. Do everything but app.exec_() -- so that we can test the application in one thread
  1656. """
  1657. app = QApplication(argv)
  1658. app.setApplicationName(__appname__)
  1659. app.setWindowIcon(newIcon("app"))
  1660. # Tzutalin 201705+: Accept extra agruments to change predefined class file
  1661. argparser = argparse.ArgumentParser()
  1662. argparser.add_argument("--lang", default='en', nargs="?")
  1663. argparser.add_argument("--predefined_classes_file",
  1664. default=os.path.join(os.path.dirname(__file__), "data", "predefined_classes.txt"),
  1665. nargs="?")
  1666. args = argparser.parse_args(argv[1:])
  1667. # Usage : labelImg.py image predefClassFile saveDir
  1668. win = MainWindow(lang=args.lang,
  1669. defaultPrefdefClassFile=args.predefined_classes_file)
  1670. win.show()
  1671. return app, win
  1672. def main():
  1673. '''construct main app and run it'''
  1674. app, _win = get_main_app(sys.argv)
  1675. return app.exec_()
  1676. if __name__ == '__main__':
  1677. resource_file = './libs/resources.py'
  1678. if not os.path.exists(resource_file):
  1679. output = os.system('pyrcc5 -o libs/resources.py resources.qrc')
  1680. assert output == 0, "operate the cmd have some problems ,please check whether there is a in the lib " \
  1681. "directory resources.py "
  1682. import libs.resources
  1683. sys.exit(main())