PPOCRLabel.py 118 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568156915701571157215731574157515761577157815791580158115821583158415851586158715881589159015911592159315941595159615971598159916001601160216031604160516061607160816091610161116121613161416151616161716181619162016211622162316241625162616271628162916301631163216331634163516361637163816391640164116421643164416451646164716481649165016511652165316541655165616571658165916601661166216631664166516661667166816691670167116721673167416751676167716781679168016811682168316841685168616871688168916901691169216931694169516961697169816991700170117021703170417051706170717081709171017111712171317141715171617171718171917201721172217231724172517261727172817291730173117321733173417351736173717381739174017411742174317441745174617471748174917501751175217531754175517561757175817591760176117621763176417651766176717681769177017711772177317741775177617771778177917801781178217831784178517861787178817891790179117921793179417951796179717981799180018011802180318041805180618071808180918101811181218131814181518161817181818191820182118221823182418251826182718281829183018311832183318341835183618371838183918401841184218431844184518461847184818491850185118521853185418551856185718581859186018611862186318641865186618671868186918701871187218731874187518761877187818791880188118821883188418851886188718881889189018911892189318941895189618971898189919001901190219031904190519061907190819091910191119121913191419151916191719181919192019211922192319241925192619271928192919301931193219331934193519361937193819391940194119421943194419451946194719481949195019511952195319541955195619571958195919601961196219631964196519661967196819691970197119721973197419751976197719781979198019811982198319841985198619871988198919901991199219931994199519961997199819992000200120022003200420052006200720082009201020112012201320142015201620172018201920202021202220232024202520262027202820292030203120322033203420352036203720382039204020412042204320442045204620472048204920502051205220532054205520562057205820592060206120622063206420652066206720682069207020712072207320742075207620772078207920802081208220832084208520862087208820892090209120922093209420952096209720982099210021012102210321042105210621072108210921102111211221132114211521162117211821192120212121222123212421252126212721282129213021312132213321342135213621372138213921402141214221432144214521462147214821492150215121522153215421552156215721582159216021612162216321642165216621672168216921702171217221732174217521762177217821792180218121822183218421852186218721882189219021912192219321942195219621972198219922002201220222032204220522062207220822092210221122122213221422152216221722182219222022212222222322242225222622272228222922302231223222332234223522362237223822392240224122422243224422452246224722482249225022512252225322542255225622572258225922602261226222632264226522662267226822692270227122722273227422752276227722782279228022812282228322842285228622872288228922902291229222932294229522962297229822992300230123022303230423052306230723082309231023112312231323142315231623172318231923202321232223232324232523262327232823292330233123322333233423352336233723382339234023412342234323442345234623472348234923502351235223532354235523562357235823592360236123622363236423652366236723682369237023712372237323742375237623772378237923802381238223832384238523862387238823892390239123922393239423952396239723982399240024012402240324042405240624072408240924102411241224132414241524162417241824192420242124222423242424252426242724282429243024312432243324342435243624372438243924402441244224432444244524462447244824492450245124522453245424552456245724582459246024612462246324642465246624672468246924702471247224732474247524762477247824792480248124822483248424852486248724882489249024912492249324942495249624972498249925002501250225032504250525062507250825092510251125122513251425152516251725182519252025212522252325242525252625272528252925302531253225332534253525362537253825392540254125422543254425452546254725482549255025512552255325542555255625572558255925602561256225632564256525662567256825692570257125722573257425752576257725782579258025812582258325842585258625872588258925902591259225932594259525962597259825992600260126022603260426052606260726082609261026112612261326142615261626172618261926202621262226232624262526262627262826292630263126322633263426352636263726382639264026412642264326442645264626472648264926502651265226532654265526562657265826592660266126622663266426652666266726682669267026712672267326742675267626772678267926802681268226832684268526862687268826892690269126922693269426952696269726982699270027012702270327042705270627072708270927102711271227132714271527162717271827192720272127222723272427252726272727282729273027312732273327342735273627372738273927402741274227432744274527462747274827492750275127522753275427552756275727582759276027612762276327642765276627672768276927702771277227732774277527762777277827792780278127822783278427852786278727882789279027912792279327942795279627972798279928002801280228032804280528062807280828092810281128122813281428152816281728182819282028212822282328242825282628272828282928302831283228332834283528362837283828392840
  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 json
  20. import os.path
  21. import platform
  22. import subprocess
  23. import sys
  24. import xlrd
  25. from functools import partial
  26. from PyQt5.QtCore import QSize, Qt, QPoint, QByteArray, QTimer, QFileInfo, QPointF, QProcess
  27. from PyQt5.QtGui import QImage, QCursor, QPixmap, QImageReader
  28. from PyQt5.QtWidgets import QMainWindow, QListWidget, QVBoxLayout, QToolButton, QHBoxLayout, QDockWidget, QWidget, \
  29. QSlider, QGraphicsOpacityEffect, QMessageBox, QListView, QScrollArea, QWidgetAction, QApplication, QLabel, QGridLayout, \
  30. QFileDialog, QListWidgetItem, QComboBox, QDialog, QAbstractItemView, QSizePolicy
  31. __dir__ = os.path.dirname(os.path.abspath(__file__))
  32. sys.path.append(__dir__)
  33. sys.path.append(os.path.abspath(os.path.join(__dir__, '../..')))
  34. sys.path.append(os.path.abspath(os.path.join(__dir__, '../PaddleOCR')))
  35. sys.path.append("..")
  36. from paddleocr import PaddleOCR, PPStructure
  37. from libs.constants import *
  38. from libs.utils import *
  39. from libs.labelColor import label_colormap
  40. from libs.settings import Settings
  41. from libs.shape import Shape, DEFAULT_LINE_COLOR, DEFAULT_FILL_COLOR, DEFAULT_LOCK_COLOR
  42. from libs.stringBundle import StringBundle
  43. from libs.canvas import Canvas
  44. from libs.zoomWidget import ZoomWidget
  45. from libs.autoDialog import AutoDialog
  46. from libs.labelDialog import LabelDialog
  47. from libs.colorDialog import ColorDialog
  48. from libs.ustr import ustr
  49. from libs.hashableQListWidgetItem import HashableQListWidgetItem
  50. from libs.editinlist import EditInList
  51. from libs.unique_label_qlist_widget import UniqueLabelQListWidget
  52. from libs.keyDialog import KeyDialog
  53. __appname__ = 'PPOCRLabel'
  54. LABEL_COLORMAP = label_colormap()
  55. class MainWindow(QMainWindow):
  56. FIT_WINDOW, FIT_WIDTH, MANUAL_ZOOM = list(range(3))
  57. def __init__(self,
  58. lang="ch",
  59. gpu=False,
  60. kie_mode=False,
  61. default_filename=None,
  62. default_predefined_class_file=None,
  63. default_save_dir=None):
  64. super(MainWindow, self).__init__()
  65. self.setWindowTitle(__appname__)
  66. self.setWindowState(Qt.WindowMaximized) # set window max
  67. self.activateWindow() # PPOCRLabel goes to the front when activate
  68. # Load setting in the main thread
  69. self.settings = Settings()
  70. self.settings.load()
  71. settings = self.settings
  72. self.lang = lang
  73. # Load string bundle for i18n
  74. if lang not in ['ch', 'en']:
  75. lang = 'en'
  76. self.stringBundle = StringBundle.getBundle(localeStr='zh-CN' if lang == 'ch' else 'en') # 'en'
  77. getStr = lambda strId: self.stringBundle.getString(strId)
  78. # KIE setting
  79. self.kie_mode = kie_mode
  80. self.key_previous_text = ""
  81. self.existed_key_cls_set = set()
  82. self.key_dialog_tip = getStr('keyDialogTip')
  83. self.defaultSaveDir = default_save_dir
  84. self.ocr = PaddleOCR(use_pdserving=False,
  85. use_angle_cls=True,
  86. det=True,
  87. cls=True,
  88. use_gpu=gpu,
  89. lang=lang,
  90. show_log=False)
  91. self.table_ocr = PPStructure(use_pdserving=False,
  92. use_gpu=gpu,
  93. lang=lang,
  94. layout=False,
  95. show_log=False)
  96. if os.path.exists('./data/paddle.png'):
  97. result = self.ocr.ocr('./data/paddle.png', cls=True, det=True)
  98. result = self.table_ocr('./data/paddle.png', return_ocr_result_in_table=True)
  99. # For loading all image under a directory
  100. self.mImgList = []
  101. self.mImgList5 = []
  102. self.dirname = None
  103. self.labelHist = []
  104. self.lastOpenDir = None
  105. self.result_dic = []
  106. self.result_dic_locked = []
  107. self.changeFileFolder = False
  108. self.haveAutoReced = False
  109. self.labelFile = None
  110. self.currIndex = 0
  111. # Whether we need to save or not.
  112. self.dirty = False
  113. self._noSelectionSlot = False
  114. self._beginner = True
  115. self.screencastViewer = self.getAvailableScreencastViewer()
  116. self.screencast = "https://github.com/PaddlePaddle/PaddleOCR"
  117. # Load predefined classes to the list
  118. self.loadPredefinedClasses(default_predefined_class_file)
  119. # Main widgets and related state.
  120. self.labelDialog = LabelDialog(parent=self, listItem=self.labelHist)
  121. self.autoDialog = AutoDialog(parent=self)
  122. self.itemsToShapes = {}
  123. self.shapesToItems = {}
  124. self.itemsToShapesbox = {}
  125. self.shapesToItemsbox = {}
  126. self.prevLabelText = getStr('tempLabel')
  127. self.noLabelText = getStr('nullLabel')
  128. self.model = 'paddle'
  129. self.PPreader = None
  130. self.autoSaveNum = 5
  131. # ================== File List ==================
  132. filelistLayout = QVBoxLayout()
  133. filelistLayout.setContentsMargins(0, 0, 0, 0)
  134. self.fileListWidget = QListWidget()
  135. self.fileListWidget.itemClicked.connect(self.fileitemDoubleClicked)
  136. self.fileListWidget.setIconSize(QSize(25, 25))
  137. filelistLayout.addWidget(self.fileListWidget)
  138. fileListContainer = QWidget()
  139. fileListContainer.setLayout(filelistLayout)
  140. self.fileListName = getStr('fileList')
  141. self.fileDock = QDockWidget(self.fileListName, self)
  142. self.fileDock.setObjectName(getStr('files'))
  143. self.fileDock.setWidget(fileListContainer)
  144. self.addDockWidget(Qt.LeftDockWidgetArea, self.fileDock)
  145. # ================== Key List ==================
  146. if self.kie_mode:
  147. self.keyList = UniqueLabelQListWidget()
  148. # set key list height
  149. key_list_height = int(QApplication.desktop().height() // 4)
  150. if key_list_height < 50:
  151. key_list_height = 50
  152. self.keyList.setMaximumHeight(key_list_height)
  153. self.keyListDockName = getStr('keyListTitle')
  154. self.keyListDock = QDockWidget(self.keyListDockName, self)
  155. self.keyListDock.setWidget(self.keyList)
  156. self.keyListDock.setFeatures(QDockWidget.NoDockWidgetFeatures)
  157. filelistLayout.addWidget(self.keyListDock)
  158. self.AutoRecognition = QToolButton()
  159. self.AutoRecognition.setToolButtonStyle(Qt.ToolButtonTextBesideIcon)
  160. self.AutoRecognition.setIcon(newIcon('Auto'))
  161. autoRecLayout = QHBoxLayout()
  162. autoRecLayout.setContentsMargins(0, 0, 0, 0)
  163. autoRecLayout.addWidget(self.AutoRecognition)
  164. autoRecContainer = QWidget()
  165. autoRecContainer.setLayout(autoRecLayout)
  166. filelistLayout.addWidget(autoRecContainer)
  167. # ================== Right Area ==================
  168. listLayout = QVBoxLayout()
  169. listLayout.setContentsMargins(0, 0, 0, 0)
  170. # Buttons
  171. self.editButton = QToolButton()
  172. self.reRecogButton = QToolButton()
  173. self.reRecogButton.setIcon(newIcon('reRec', 30))
  174. self.reRecogButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon)
  175. self.tableRecButton = QToolButton()
  176. self.tableRecButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon)
  177. self.newButton = QToolButton()
  178. self.newButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon)
  179. self.createpolyButton = QToolButton()
  180. self.createpolyButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon)
  181. self.SaveButton = QToolButton()
  182. self.SaveButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon)
  183. self.DelButton = QToolButton()
  184. self.DelButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon)
  185. leftTopToolBox = QGridLayout()
  186. leftTopToolBox.addWidget(self.newButton, 0, 0, 1, 1)
  187. leftTopToolBox.addWidget(self.createpolyButton, 0, 1, 1, 1)
  188. leftTopToolBox.addWidget(self.reRecogButton, 1, 0, 1, 1)
  189. leftTopToolBox.addWidget(self.tableRecButton, 1, 1, 1, 1)
  190. leftTopToolBoxContainer = QWidget()
  191. leftTopToolBoxContainer.setLayout(leftTopToolBox)
  192. listLayout.addWidget(leftTopToolBoxContainer)
  193. # ================== Label List ==================
  194. labelIndexListlBox = QHBoxLayout()
  195. # Create and add a widget for showing current label item index
  196. self.indexList = QListWidget()
  197. self.indexList.setMaximumSize(30, 16777215) # limit max width
  198. self.indexList.setEditTriggers(QAbstractItemView.NoEditTriggers) # no editable
  199. self.indexList.itemSelectionChanged.connect(self.indexSelectionChanged)
  200. self.indexList.setVerticalScrollBarPolicy(Qt.ScrollBarAlwaysOff) # no scroll Bar
  201. self.indexListDock = QDockWidget('No.', self)
  202. self.indexListDock.setWidget(self.indexList)
  203. self.indexListDock.setFeatures(QDockWidget.NoDockWidgetFeatures)
  204. labelIndexListlBox.addWidget(self.indexListDock, 1)
  205. # no margin between two boxes
  206. labelIndexListlBox.setSpacing(0)
  207. # Create and add a widget for showing current label items
  208. self.labelList = EditInList()
  209. labelListContainer = QWidget()
  210. labelListContainer.setLayout(listLayout)
  211. self.labelList.itemSelectionChanged.connect(self.labelSelectionChanged)
  212. self.labelList.clicked.connect(self.labelList.item_clicked)
  213. # Connect to itemChanged to detect checkbox changes.
  214. self.labelList.itemChanged.connect(self.labelItemChanged)
  215. self.labelListDockName = getStr('recognitionResult')
  216. self.labelListDock = QDockWidget(self.labelListDockName, self)
  217. self.labelListDock.setWidget(self.labelList)
  218. self.labelListDock.setFeatures(QDockWidget.NoDockWidgetFeatures)
  219. labelIndexListlBox.addWidget(self.labelListDock, 10) # label list is wider than index list
  220. # enable labelList drag_drop to adjust bbox order
  221. # 设置选择模式为单选
  222. self.labelList.setSelectionMode(QAbstractItemView.SingleSelection)
  223. # 启用拖拽
  224. self.labelList.setDragEnabled(True)
  225. # 设置接受拖放
  226. self.labelList.viewport().setAcceptDrops(True)
  227. # 设置显示将要被放置的位置
  228. self.labelList.setDropIndicatorShown(True)
  229. # 设置拖放模式为移动项目,如果不设置,默认为复制项目
  230. self.labelList.setDragDropMode(QAbstractItemView.InternalMove)
  231. # 触发放置
  232. self.labelList.model().rowsMoved.connect(self.drag_drop_happened)
  233. labelIndexListContainer = QWidget()
  234. labelIndexListContainer.setLayout(labelIndexListlBox)
  235. listLayout.addWidget(labelIndexListContainer)
  236. # labelList indexList同步滚动
  237. self.labelListBar = self.labelList.verticalScrollBar()
  238. self.indexListBar = self.indexList.verticalScrollBar()
  239. self.labelListBar.valueChanged.connect(self.move_scrollbar)
  240. self.indexListBar.valueChanged.connect(self.move_scrollbar)
  241. # ================== Detection Box ==================
  242. self.BoxList = QListWidget()
  243. # self.BoxList.itemActivated.connect(self.boxSelectionChanged)
  244. self.BoxList.itemSelectionChanged.connect(self.boxSelectionChanged)
  245. self.BoxList.itemDoubleClicked.connect(self.editBox)
  246. # Connect to itemChanged to detect checkbox changes.
  247. self.BoxList.itemChanged.connect(self.boxItemChanged)
  248. self.BoxListDockName = getStr('detectionBoxposition')
  249. self.BoxListDock = QDockWidget(self.BoxListDockName, self)
  250. self.BoxListDock.setWidget(self.BoxList)
  251. self.BoxListDock.setFeatures(QDockWidget.NoDockWidgetFeatures)
  252. listLayout.addWidget(self.BoxListDock)
  253. # ================== Lower Right Area ==================
  254. leftbtmtoolbox = QHBoxLayout()
  255. leftbtmtoolbox.addWidget(self.SaveButton)
  256. leftbtmtoolbox.addWidget(self.DelButton)
  257. leftbtmtoolboxcontainer = QWidget()
  258. leftbtmtoolboxcontainer.setLayout(leftbtmtoolbox)
  259. listLayout.addWidget(leftbtmtoolboxcontainer)
  260. self.dock = QDockWidget(getStr('boxLabelText'), self)
  261. self.dock.setObjectName(getStr('labels'))
  262. self.dock.setWidget(labelListContainer)
  263. # ================== Zoom Bar ==================
  264. self.imageSlider = QSlider(Qt.Horizontal)
  265. self.imageSlider.valueChanged.connect(self.CanvasSizeChange)
  266. self.imageSlider.setMinimum(-9)
  267. self.imageSlider.setMaximum(510)
  268. self.imageSlider.setSingleStep(1)
  269. self.imageSlider.setTickPosition(QSlider.TicksBelow)
  270. self.imageSlider.setTickInterval(1)
  271. op = QGraphicsOpacityEffect()
  272. op.setOpacity(0.2)
  273. self.imageSlider.setGraphicsEffect(op)
  274. self.imageSlider.setStyleSheet("background-color:transparent")
  275. self.imageSliderDock = QDockWidget(getStr('ImageResize'), self)
  276. self.imageSliderDock.setObjectName(getStr('IR'))
  277. self.imageSliderDock.setWidget(self.imageSlider)
  278. self.imageSliderDock.setFeatures(QDockWidget.DockWidgetFloatable)
  279. self.imageSliderDock.setAttribute(Qt.WA_TranslucentBackground)
  280. self.addDockWidget(Qt.RightDockWidgetArea, self.imageSliderDock)
  281. self.zoomWidget = ZoomWidget()
  282. self.colorDialog = ColorDialog(parent=self)
  283. self.zoomWidgetValue = self.zoomWidget.value()
  284. self.msgBox = QMessageBox()
  285. # ================== Thumbnail ==================
  286. hlayout = QHBoxLayout()
  287. m = (0, 0, 0, 0)
  288. hlayout.setSpacing(0)
  289. hlayout.setContentsMargins(*m)
  290. self.preButton = QToolButton()
  291. self.preButton.setIcon(newIcon("prev", 40))
  292. self.preButton.setIconSize(QSize(40, 100))
  293. self.preButton.clicked.connect(self.openPrevImg)
  294. self.preButton.setStyleSheet('border: none;')
  295. self.preButton.setShortcut('a')
  296. self.iconlist = QListWidget()
  297. self.iconlist.setViewMode(QListView.IconMode)
  298. self.iconlist.setFlow(QListView.TopToBottom)
  299. self.iconlist.setSpacing(10)
  300. self.iconlist.setIconSize(QSize(50, 50))
  301. self.iconlist.setMovement(QListView.Static)
  302. self.iconlist.setResizeMode(QListView.Adjust)
  303. self.iconlist.itemClicked.connect(self.iconitemDoubleClicked)
  304. self.iconlist.setStyleSheet("QListWidget{ background-color:transparent; border: none;}")
  305. self.iconlist.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff)
  306. self.nextButton = QToolButton()
  307. self.nextButton.setIcon(newIcon("next", 40))
  308. self.nextButton.setIconSize(QSize(40, 100))
  309. self.nextButton.setStyleSheet('border: none;')
  310. self.nextButton.clicked.connect(self.openNextImg)
  311. self.nextButton.setShortcut('d')
  312. hlayout.addWidget(self.preButton)
  313. hlayout.addWidget(self.iconlist)
  314. hlayout.addWidget(self.nextButton)
  315. iconListContainer = QWidget()
  316. iconListContainer.setLayout(hlayout)
  317. iconListContainer.setFixedHeight(100)
  318. # ================== Canvas ==================
  319. self.canvas = Canvas(parent=self)
  320. self.canvas.zoomRequest.connect(self.zoomRequest)
  321. self.canvas.setDrawingShapeToSquare(settings.get(SETTING_DRAW_SQUARE, False))
  322. scroll = QScrollArea()
  323. scroll.setWidget(self.canvas)
  324. scroll.setWidgetResizable(True)
  325. self.scrollBars = {
  326. Qt.Vertical: scroll.verticalScrollBar(),
  327. Qt.Horizontal: scroll.horizontalScrollBar()
  328. }
  329. self.scrollArea = scroll
  330. self.canvas.scrollRequest.connect(self.scrollRequest)
  331. self.canvas.newShape.connect(partial(self.newShape, False))
  332. self.canvas.shapeMoved.connect(self.updateBoxlist) # self.setDirty
  333. self.canvas.selectionChanged.connect(self.shapeSelectionChanged)
  334. self.canvas.drawingPolygon.connect(self.toggleDrawingSensitive)
  335. centerLayout = QVBoxLayout()
  336. centerLayout.setContentsMargins(0, 0, 0, 0)
  337. centerLayout.addWidget(scroll)
  338. centerLayout.addWidget(iconListContainer, 0, Qt.AlignCenter)
  339. centerContainer = QWidget()
  340. centerContainer.setLayout(centerLayout)
  341. self.setCentralWidget(centerContainer)
  342. self.addDockWidget(Qt.RightDockWidgetArea, self.dock)
  343. self.dock.setFeatures(QDockWidget.DockWidgetClosable | QDockWidget.DockWidgetFloatable)
  344. self.fileDock.setFeatures(QDockWidget.NoDockWidgetFeatures)
  345. # ================== Actions ==================
  346. action = partial(newAction, self)
  347. quit = action(getStr('quit'), self.close,
  348. 'Ctrl+Q', 'quit', getStr('quitApp'))
  349. opendir = action(getStr('openDir'), self.openDirDialog,
  350. 'Ctrl+u', 'open', getStr('openDir'))
  351. open_dataset_dir = action(getStr('openDatasetDir'), self.openDatasetDirDialog,
  352. 'Ctrl+p', 'open', getStr('openDatasetDir'), enabled=False)
  353. save = action(getStr('save'), self.saveFile,
  354. 'Ctrl+V', 'verify', getStr('saveDetail'), enabled=False)
  355. alcm = action(getStr('choosemodel'), self.autolcm,
  356. 'Ctrl+M', 'next', getStr('tipchoosemodel'))
  357. deleteImg = action(getStr('deleteImg'), self.deleteImg, 'Ctrl+Shift+D', 'close', getStr('deleteImgDetail'),
  358. enabled=True)
  359. resetAll = action(getStr('resetAll'), self.resetAll, None, 'resetall', getStr('resetAllDetail'))
  360. color1 = action(getStr('boxLineColor'), self.chooseColor,
  361. 'Ctrl+L', 'color_line', getStr('boxLineColorDetail'))
  362. createMode = action(getStr('crtBox'), self.setCreateMode,
  363. 'w', 'new', getStr('crtBoxDetail'), enabled=False)
  364. editMode = action('&Edit\nRectBox', self.setEditMode,
  365. 'Ctrl+J', 'edit', u'Move and edit Boxs', enabled=False)
  366. create = action(getStr('crtBox'), self.createShape,
  367. 'w', 'objects', getStr('crtBoxDetail'), enabled=False)
  368. delete = action(getStr('delBox'), self.deleteSelectedShape,
  369. 'backspace', 'delete', getStr('delBoxDetail'), enabled=False)
  370. copy = action(getStr('dupBox'), self.copySelectedShape,
  371. 'Ctrl+C', 'copy', getStr('dupBoxDetail'),
  372. enabled=False)
  373. hideAll = action(getStr('hideBox'), partial(self.togglePolygons, False),
  374. 'Ctrl+H', 'hide', getStr('hideAllBoxDetail'),
  375. enabled=False)
  376. showAll = action(getStr('showBox'), partial(self.togglePolygons, True),
  377. 'Ctrl+A', 'hide', getStr('showAllBoxDetail'),
  378. enabled=False)
  379. help = action(getStr('tutorial'), self.showTutorialDialog, None, 'help', getStr('tutorialDetail'))
  380. showInfo = action(getStr('info'), self.showInfoDialog, None, 'help', getStr('info'))
  381. showSteps = action(getStr('steps'), self.showStepsDialog, None, 'help', getStr('steps'))
  382. showKeys = action(getStr('keys'), self.showKeysDialog, None, 'help', getStr('keys'))
  383. zoom = QWidgetAction(self)
  384. zoom.setDefaultWidget(self.zoomWidget)
  385. self.zoomWidget.setWhatsThis(
  386. u"Zoom in or out of the image. Also accessible with"
  387. " %s and %s from the canvas." % (fmtShortcut("Ctrl+[-+]"),
  388. fmtShortcut("Ctrl+Wheel")))
  389. self.zoomWidget.setEnabled(False)
  390. zoomIn = action(getStr('zoomin'), partial(self.addZoom, 10),
  391. 'Ctrl++', 'zoom-in', getStr('zoominDetail'), enabled=False)
  392. zoomOut = action(getStr('zoomout'), partial(self.addZoom, -10),
  393. 'Ctrl+-', 'zoom-out', getStr('zoomoutDetail'), enabled=False)
  394. zoomOrg = action(getStr('originalsize'), partial(self.setZoom, 100),
  395. 'Ctrl+=', 'zoom', getStr('originalsizeDetail'), enabled=False)
  396. fitWindow = action(getStr('fitWin'), self.setFitWindow,
  397. 'Ctrl+F', 'fit-window', getStr('fitWinDetail'),
  398. checkable=True, enabled=False)
  399. fitWidth = action(getStr('fitWidth'), self.setFitWidth,
  400. 'Ctrl+Shift+F', 'fit-width', getStr('fitWidthDetail'),
  401. checkable=True, enabled=False)
  402. # Group zoom controls into a list for easier toggling.
  403. zoomActions = (self.zoomWidget, zoomIn, zoomOut,
  404. zoomOrg, fitWindow, fitWidth)
  405. self.zoomMode = self.MANUAL_ZOOM
  406. self.scalers = {
  407. self.FIT_WINDOW: self.scaleFitWindow,
  408. self.FIT_WIDTH: self.scaleFitWidth,
  409. # Set to one to scale to 100% when loading files.
  410. self.MANUAL_ZOOM: lambda: 1,
  411. }
  412. # ================== New Actions ==================
  413. edit = action(getStr('editLabel'), self.editLabel,
  414. 'Ctrl+E', 'edit', getStr('editLabelDetail'), enabled=False)
  415. AutoRec = action(getStr('autoRecognition'), self.autoRecognition,
  416. '', 'Auto', getStr('autoRecognition'), enabled=False)
  417. reRec = action(getStr('reRecognition'), self.reRecognition,
  418. 'Ctrl+Shift+R', 'reRec', getStr('reRecognition'), enabled=False)
  419. singleRere = action(getStr('singleRe'), self.singleRerecognition,
  420. 'Ctrl+R', 'reRec', getStr('singleRe'), enabled=False)
  421. createpoly = action(getStr('creatPolygon'), self.createPolygon,
  422. 'q', 'new', getStr('creatPolygon'), enabled=False)
  423. tableRec = action(getStr('TableRecognition'), self.TableRecognition,
  424. '', 'Auto', getStr('TableRecognition'), enabled=False)
  425. cellreRec = action(getStr('cellreRecognition'), self.cellreRecognition,
  426. '', 'reRec', getStr('cellreRecognition'), enabled=False)
  427. saveRec = action(getStr('saveRec'), self.saveRecResult,
  428. '', 'save', getStr('saveRec'), enabled=False)
  429. saveLabel = action(getStr('saveLabel'), self.saveLabelFile, #
  430. 'Ctrl+S', 'save', getStr('saveLabel'), enabled=False)
  431. exportJSON = action(getStr('exportJSON'), self.exportJSON,
  432. '', 'save', getStr('exportJSON'), enabled=False)
  433. undoLastPoint = action(getStr("undoLastPoint"), self.canvas.undoLastPoint,
  434. 'Ctrl+Z', "undo", getStr("undoLastPoint"), enabled=False)
  435. rotateLeft = action(getStr("rotateLeft"), partial(self.rotateImgAction, 1),
  436. 'Ctrl+Alt+L', "rotateLeft", getStr("rotateLeft"), enabled=False)
  437. rotateRight = action(getStr("rotateRight"), partial(self.rotateImgAction, -1),
  438. 'Ctrl+Alt+R', "rotateRight", getStr("rotateRight"), enabled=False)
  439. undo = action(getStr("undo"), self.undoShapeEdit,
  440. 'Ctrl+Z', "undo", getStr("undo"), enabled=False)
  441. change_cls = action(getStr("keyChange"), self.change_box_key,
  442. 'Ctrl+X', "edit", getStr("keyChange"), enabled=False)
  443. lock = action(getStr("lockBox"), self.lockSelectedShape,
  444. None, "lock", getStr("lockBoxDetail"), enabled=False)
  445. self.editButton.setDefaultAction(edit)
  446. self.newButton.setDefaultAction(create)
  447. self.createpolyButton.setDefaultAction(createpoly)
  448. self.DelButton.setDefaultAction(deleteImg)
  449. self.SaveButton.setDefaultAction(save)
  450. self.AutoRecognition.setDefaultAction(AutoRec)
  451. self.reRecogButton.setDefaultAction(reRec)
  452. self.tableRecButton.setDefaultAction(tableRec)
  453. # self.preButton.setDefaultAction(openPrevImg)
  454. # self.nextButton.setDefaultAction(openNextImg)
  455. # ================== Zoom layout ==================
  456. zoomLayout = QHBoxLayout()
  457. zoomLayout.addStretch()
  458. self.zoominButton = QToolButton()
  459. self.zoominButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon)
  460. self.zoominButton.setDefaultAction(zoomIn)
  461. self.zoomoutButton = QToolButton()
  462. self.zoomoutButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon)
  463. self.zoomoutButton.setDefaultAction(zoomOut)
  464. self.zoomorgButton = QToolButton()
  465. self.zoomorgButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon)
  466. self.zoomorgButton.setDefaultAction(zoomOrg)
  467. zoomLayout.addWidget(self.zoominButton)
  468. zoomLayout.addWidget(self.zoomorgButton)
  469. zoomLayout.addWidget(self.zoomoutButton)
  470. zoomContainer = QWidget()
  471. zoomContainer.setLayout(zoomLayout)
  472. zoomContainer.setGeometry(0, 0, 30, 150)
  473. shapeLineColor = action(getStr('shapeLineColor'), self.chshapeLineColor,
  474. icon='color_line', tip=getStr('shapeLineColorDetail'),
  475. enabled=False)
  476. shapeFillColor = action(getStr('shapeFillColor'), self.chshapeFillColor,
  477. icon='color', tip=getStr('shapeFillColorDetail'),
  478. enabled=False)
  479. # Label list context menu.
  480. labelMenu = QMenu()
  481. addActions(labelMenu, (edit, delete))
  482. self.labelList.setContextMenuPolicy(Qt.CustomContextMenu)
  483. self.labelList.customContextMenuRequested.connect(self.popLabelListMenu)
  484. # Draw squares/rectangles
  485. self.drawSquaresOption = QAction(getStr('drawSquares'), self)
  486. self.drawSquaresOption.setCheckable(True)
  487. self.drawSquaresOption.setChecked(settings.get(SETTING_DRAW_SQUARE, False))
  488. self.drawSquaresOption.triggered.connect(self.toogleDrawSquare)
  489. # Store actions for further handling.
  490. self.actions = struct(save=save, resetAll=resetAll, deleteImg=deleteImg,
  491. lineColor=color1, create=create, createpoly=createpoly, tableRec=tableRec, delete=delete, edit=edit, copy=copy,
  492. saveRec=saveRec, singleRere=singleRere, AutoRec=AutoRec, reRec=reRec, cellreRec=cellreRec,
  493. createMode=createMode, editMode=editMode,
  494. shapeLineColor=shapeLineColor, shapeFillColor=shapeFillColor,
  495. zoom=zoom, zoomIn=zoomIn, zoomOut=zoomOut, zoomOrg=zoomOrg,
  496. fitWindow=fitWindow, fitWidth=fitWidth,
  497. zoomActions=zoomActions, saveLabel=saveLabel, change_cls=change_cls,
  498. undo=undo, undoLastPoint=undoLastPoint, open_dataset_dir=open_dataset_dir,
  499. rotateLeft=rotateLeft, rotateRight=rotateRight, lock=lock, exportJSON=exportJSON,
  500. fileMenuActions=(opendir, open_dataset_dir, saveLabel, exportJSON, resetAll, quit),
  501. beginner=(), advanced=(),
  502. editMenu=(createpoly, edit, copy, delete, singleRere, cellreRec, None, undo, undoLastPoint,
  503. None, rotateLeft, rotateRight, None, color1, self.drawSquaresOption, lock,
  504. None, change_cls),
  505. beginnerContext=(
  506. create, createpoly, edit, copy, delete, singleRere, cellreRec, rotateLeft, rotateRight, lock, change_cls),
  507. advancedContext=(createMode, editMode, edit, copy,
  508. delete, shapeLineColor, shapeFillColor),
  509. onLoadActive=(create, createpoly, createMode, editMode),
  510. onShapesPresent=(hideAll, showAll))
  511. # menus
  512. self.menus = struct(
  513. file=self.menu('&' + getStr('mfile')),
  514. edit=self.menu('&' + getStr('medit')),
  515. view=self.menu('&' + getStr('mview')),
  516. autolabel=self.menu('&PaddleOCR'),
  517. help=self.menu('&' + getStr('mhelp')),
  518. recentFiles=QMenu('Open &Recent'),
  519. labelList=labelMenu)
  520. self.lastLabel = None
  521. # Add option to enable/disable labels being displayed at the top of bounding boxes
  522. self.displayLabelOption = QAction(getStr('displayLabel'), self)
  523. self.displayLabelOption.setShortcut("Ctrl+Shift+P")
  524. self.displayLabelOption.setCheckable(True)
  525. self.displayLabelOption.setChecked(settings.get(SETTING_PAINT_LABEL, False))
  526. self.displayLabelOption.triggered.connect(self.togglePaintLabelsOption)
  527. # Add option to enable/disable box index being displayed at the top of bounding boxes
  528. self.displayIndexOption = QAction(getStr('displayIndex'), self)
  529. self.displayIndexOption.setCheckable(True)
  530. self.displayIndexOption.setChecked(settings.get(SETTING_PAINT_INDEX, False))
  531. self.displayIndexOption.triggered.connect(self.togglePaintIndexOption)
  532. self.labelDialogOption = QAction(getStr('labelDialogOption'), self)
  533. self.labelDialogOption.setShortcut("Ctrl+Shift+L")
  534. self.labelDialogOption.setCheckable(True)
  535. self.labelDialogOption.setChecked(settings.get(SETTING_PAINT_LABEL, False))
  536. self.displayIndexOption.setChecked(settings.get(SETTING_PAINT_INDEX, False))
  537. self.labelDialogOption.triggered.connect(self.speedChoose)
  538. self.autoSaveOption = QAction(getStr('autoSaveMode'), self)
  539. self.autoSaveOption.setCheckable(True)
  540. self.autoSaveOption.setChecked(settings.get(SETTING_PAINT_LABEL, False))
  541. self.displayIndexOption.setChecked(settings.get(SETTING_PAINT_INDEX, False))
  542. self.autoSaveOption.triggered.connect(self.autoSaveFunc)
  543. addActions(self.menus.file,
  544. (opendir, open_dataset_dir, None, saveLabel, saveRec, exportJSON, self.autoSaveOption, None, resetAll, deleteImg,
  545. quit))
  546. addActions(self.menus.help, (showKeys, showSteps, showInfo))
  547. addActions(self.menus.view, (
  548. self.displayLabelOption, self.displayIndexOption, self.labelDialogOption,
  549. None,
  550. hideAll, showAll, None,
  551. zoomIn, zoomOut, zoomOrg, None,
  552. fitWindow, fitWidth))
  553. addActions(self.menus.autolabel, (AutoRec, reRec, cellreRec, alcm, None, help))
  554. self.menus.file.aboutToShow.connect(self.updateFileMenu)
  555. # Custom context menu for the canvas widget:
  556. addActions(self.canvas.menus[0], self.actions.beginnerContext)
  557. self.statusBar().showMessage('%s started.' % __appname__)
  558. self.statusBar().show()
  559. # Application state.
  560. self.image = QImage()
  561. self.filePath = ustr(default_filename)
  562. self.lastOpenDir = None
  563. self.recentFiles = []
  564. self.maxRecent = 7
  565. self.lineColor = None
  566. self.fillColor = None
  567. self.zoom_level = 100
  568. self.fit_window = False
  569. # Add Chris
  570. self.difficult = False
  571. # Fix the compatible issue for qt4 and qt5. Convert the QStringList to python list
  572. if settings.get(SETTING_RECENT_FILES):
  573. if have_qstring():
  574. recentFileQStringList = settings.get(SETTING_RECENT_FILES)
  575. self.recentFiles = [ustr(i) for i in recentFileQStringList]
  576. else:
  577. self.recentFiles = recentFileQStringList = settings.get(SETTING_RECENT_FILES)
  578. size = settings.get(SETTING_WIN_SIZE, QSize(1200, 800))
  579. position = QPoint(0, 0)
  580. saved_position = settings.get(SETTING_WIN_POSE, position)
  581. # Fix the multiple monitors issue
  582. for i in range(QApplication.desktop().screenCount()):
  583. if QApplication.desktop().availableGeometry(i).contains(saved_position):
  584. position = saved_position
  585. break
  586. self.resize(size)
  587. self.move(position)
  588. saveDir = ustr(settings.get(SETTING_SAVE_DIR, None))
  589. self.lastOpenDir = ustr(settings.get(SETTING_LAST_OPEN_DIR, None))
  590. self.restoreState(settings.get(SETTING_WIN_STATE, QByteArray()))
  591. Shape.line_color = self.lineColor = QColor(settings.get(SETTING_LINE_COLOR, DEFAULT_LINE_COLOR))
  592. Shape.fill_color = self.fillColor = QColor(settings.get(SETTING_FILL_COLOR, DEFAULT_FILL_COLOR))
  593. self.canvas.setDrawingColor(self.lineColor)
  594. # Add chris
  595. Shape.difficult = self.difficult
  596. # ADD:
  597. # Populate the File menu dynamically.
  598. self.updateFileMenu()
  599. # Since loading the file may take some time, make sure it runs in the background.
  600. if self.filePath and os.path.isdir(self.filePath):
  601. self.queueEvent(partial(self.importDirImages, self.filePath or ""))
  602. elif self.filePath:
  603. self.queueEvent(partial(self.loadFile, self.filePath or ""))
  604. self.keyDialog = None
  605. # Callbacks:
  606. self.zoomWidget.valueChanged.connect(self.paintCanvas)
  607. self.populateModeActions()
  608. # Display cursor coordinates at the right of status bar
  609. self.labelCoordinates = QLabel('')
  610. self.statusBar().addPermanentWidget(self.labelCoordinates)
  611. # Open Dir if deafult file
  612. if self.filePath and os.path.isdir(self.filePath):
  613. self.openDirDialog(dirpath=self.filePath, silent=True)
  614. def menu(self, title, actions=None):
  615. menu = self.menuBar().addMenu(title)
  616. if actions:
  617. addActions(menu, actions)
  618. return menu
  619. def keyReleaseEvent(self, event):
  620. if event.key() == Qt.Key_Control:
  621. self.canvas.setDrawingShapeToSquare(False)
  622. def keyPressEvent(self, event):
  623. if event.key() == Qt.Key_Control:
  624. # Draw rectangle if Ctrl is pressed
  625. self.canvas.setDrawingShapeToSquare(True)
  626. def noShapes(self):
  627. return not self.itemsToShapes
  628. def populateModeActions(self):
  629. self.canvas.menus[0].clear()
  630. addActions(self.canvas.menus[0], self.actions.beginnerContext)
  631. self.menus.edit.clear()
  632. actions = (self.actions.create,) # if self.beginner() else (self.actions.createMode, self.actions.editMode)
  633. addActions(self.menus.edit, actions + self.actions.editMenu)
  634. def setDirty(self):
  635. self.dirty = True
  636. self.actions.save.setEnabled(True)
  637. def setClean(self):
  638. self.dirty = False
  639. self.actions.save.setEnabled(False)
  640. self.actions.create.setEnabled(True)
  641. self.actions.createpoly.setEnabled(True)
  642. def toggleActions(self, value=True):
  643. """Enable/Disable widgets which depend on an opened image."""
  644. for z in self.actions.zoomActions:
  645. z.setEnabled(value)
  646. for action in self.actions.onLoadActive:
  647. action.setEnabled(value)
  648. def queueEvent(self, function):
  649. QTimer.singleShot(0, function)
  650. def status(self, message, delay=5000):
  651. self.statusBar().showMessage(message, delay)
  652. def resetState(self):
  653. self.itemsToShapes.clear()
  654. self.shapesToItems.clear()
  655. self.itemsToShapesbox.clear() # ADD
  656. self.shapesToItemsbox.clear()
  657. self.labelList.clear()
  658. self.BoxList.clear()
  659. self.indexList.clear()
  660. self.filePath = None
  661. self.imageData = None
  662. self.labelFile = None
  663. self.canvas.resetState()
  664. self.labelCoordinates.clear()
  665. # self.comboBox.cb.clear()
  666. self.result_dic = []
  667. def currentItem(self):
  668. items = self.labelList.selectedItems()
  669. if items:
  670. return items[0]
  671. return None
  672. def currentBox(self):
  673. items = self.BoxList.selectedItems()
  674. if items:
  675. return items[0]
  676. return None
  677. def addRecentFile(self, filePath):
  678. if filePath in self.recentFiles:
  679. self.recentFiles.remove(filePath)
  680. elif len(self.recentFiles) >= self.maxRecent:
  681. self.recentFiles.pop()
  682. self.recentFiles.insert(0, filePath)
  683. def beginner(self):
  684. return self._beginner
  685. def advanced(self):
  686. return not self.beginner()
  687. def getAvailableScreencastViewer(self):
  688. osName = platform.system()
  689. if osName == 'Windows':
  690. return ['C:\\Program Files\\Internet Explorer\\iexplore.exe']
  691. elif osName == 'Linux':
  692. return ['xdg-open']
  693. elif osName == 'Darwin':
  694. return ['open']
  695. ## Callbacks ##
  696. def showTutorialDialog(self):
  697. subprocess.Popen(self.screencastViewer + [self.screencast])
  698. def showInfoDialog(self):
  699. from libs.__init__ import __version__
  700. msg = u'Name:{0} \nApp Version:{1} \n{2} '.format(__appname__, __version__, sys.version_info)
  701. QMessageBox.information(self, u'Information', msg)
  702. def showStepsDialog(self):
  703. msg = stepsInfo(self.lang)
  704. QMessageBox.information(self, u'Information', msg)
  705. def showKeysDialog(self):
  706. msg = keysInfo(self.lang)
  707. QMessageBox.information(self, u'Information', msg)
  708. def createShape(self):
  709. assert self.beginner()
  710. self.canvas.setEditing(False)
  711. self.actions.create.setEnabled(False)
  712. self.actions.createpoly.setEnabled(False)
  713. self.canvas.fourpoint = False
  714. def createPolygon(self):
  715. assert self.beginner()
  716. self.canvas.setEditing(False)
  717. self.canvas.fourpoint = True
  718. self.actions.create.setEnabled(False)
  719. self.actions.createpoly.setEnabled(False)
  720. self.actions.undoLastPoint.setEnabled(True)
  721. def rotateImg(self, filename, k, _value):
  722. self.actions.rotateRight.setEnabled(_value)
  723. pix = cv2.imread(filename)
  724. pix = np.rot90(pix, k)
  725. cv2.imwrite(filename, pix)
  726. self.canvas.update()
  727. self.loadFile(filename)
  728. def rotateImgWarn(self):
  729. if self.lang == 'ch':
  730. self.msgBox.warning(self, "提示", "\n 该图片已经有标注框,旋转操作会打乱标注,建议清除标注框后旋转。")
  731. else:
  732. self.msgBox.warning(self, "Warn", "\n The picture already has a label box, "
  733. "and rotation will disrupt the label. "
  734. "It is recommended to clear the label box and rotate it.")
  735. def rotateImgAction(self, k=1, _value=False):
  736. filename = self.mImgList[self.currIndex]
  737. if os.path.exists(filename):
  738. if self.itemsToShapesbox:
  739. self.rotateImgWarn()
  740. else:
  741. self.saveFile()
  742. self.dirty = False
  743. self.rotateImg(filename=filename, k=k, _value=True)
  744. else:
  745. self.rotateImgWarn()
  746. self.actions.rotateRight.setEnabled(False)
  747. self.actions.rotateLeft.setEnabled(False)
  748. def toggleDrawingSensitive(self, drawing=True):
  749. """In the middle of drawing, toggling between modes should be disabled."""
  750. self.actions.editMode.setEnabled(not drawing)
  751. if not drawing and self.beginner():
  752. # Cancel creation.
  753. print('Cancel creation.')
  754. self.canvas.setEditing(True)
  755. self.canvas.restoreCursor()
  756. self.actions.create.setEnabled(True)
  757. self.actions.createpoly.setEnabled(True)
  758. def toggleDrawMode(self, edit=True):
  759. self.canvas.setEditing(edit)
  760. self.actions.createMode.setEnabled(edit)
  761. self.actions.editMode.setEnabled(not edit)
  762. def setCreateMode(self):
  763. assert self.advanced()
  764. self.toggleDrawMode(False)
  765. def setEditMode(self):
  766. assert self.advanced()
  767. self.toggleDrawMode(True)
  768. self.labelSelectionChanged()
  769. def updateFileMenu(self):
  770. currFilePath = self.filePath
  771. def exists(filename):
  772. return os.path.exists(filename)
  773. menu = self.menus.recentFiles
  774. menu.clear()
  775. files = [f for f in self.recentFiles if f !=
  776. currFilePath and exists(f)]
  777. for i, f in enumerate(files):
  778. icon = newIcon('labels')
  779. action = QAction(
  780. icon, '&%d %s' % (i + 1, QFileInfo(f).fileName()), self)
  781. action.triggered.connect(partial(self.loadRecent, f))
  782. menu.addAction(action)
  783. def popLabelListMenu(self, point):
  784. self.menus.labelList.exec_(self.labelList.mapToGlobal(point))
  785. def editLabel(self):
  786. if not self.canvas.editing():
  787. return
  788. item = self.currentItem()
  789. if not item:
  790. return
  791. text = self.labelDialog.popUp(item.text())
  792. if text is not None:
  793. item.setText(text)
  794. # item.setBackground(generateColorByText(text))
  795. self.setDirty()
  796. self.updateComboBox()
  797. # =================== detection box related functions ===================
  798. def boxItemChanged(self, item):
  799. shape = self.itemsToShapesbox[item]
  800. box = ast.literal_eval(item.text())
  801. # print('shape in labelItemChanged is',shape.points)
  802. if box != [(int(p.x()), int(p.y())) for p in shape.points]:
  803. # shape.points = box
  804. shape.points = [QPointF(p[0], p[1]) for p in box]
  805. # QPointF(x,y)
  806. # shape.line_color = generateColorByText(shape.label)
  807. self.setDirty()
  808. else: # User probably changed item visibility
  809. self.canvas.setShapeVisible(shape, True) # item.checkState() == Qt.Checked
  810. def editBox(self): # ADD
  811. if not self.canvas.editing():
  812. return
  813. item = self.currentBox()
  814. if not item:
  815. return
  816. text = self.labelDialog.popUp(item.text())
  817. imageSize = str(self.image.size())
  818. width, height = self.image.width(), self.image.height()
  819. if text:
  820. try:
  821. text_list = eval(text)
  822. except:
  823. msg_box = QMessageBox(QMessageBox.Warning, 'Warning', 'Please enter the correct format')
  824. msg_box.exec_()
  825. return
  826. if len(text_list) < 4:
  827. msg_box = QMessageBox(QMessageBox.Warning, 'Warning', 'Please enter the coordinates of 4 points')
  828. msg_box.exec_()
  829. return
  830. for box in text_list:
  831. if box[0] > width or box[0] < 0 or box[1] > height or box[1] < 0:
  832. msg_box = QMessageBox(QMessageBox.Warning, 'Warning', 'Out of picture size')
  833. msg_box.exec_()
  834. return
  835. item.setText(text)
  836. # item.setBackground(generateColorByText(text))
  837. self.setDirty()
  838. self.updateComboBox()
  839. def updateBoxlist(self):
  840. self.canvas.selectedShapes_hShape = []
  841. if self.canvas.hShape != None:
  842. self.canvas.selectedShapes_hShape = self.canvas.selectedShapes + [self.canvas.hShape]
  843. else:
  844. self.canvas.selectedShapes_hShape = self.canvas.selectedShapes
  845. for shape in self.canvas.selectedShapes_hShape:
  846. if shape in self.shapesToItemsbox.keys():
  847. item = self.shapesToItemsbox[shape] # listitem
  848. text = [(int(p.x()), int(p.y())) for p in shape.points]
  849. item.setText(str(text))
  850. self.actions.undo.setEnabled(True)
  851. self.setDirty()
  852. def indexTo5Files(self, currIndex):
  853. if currIndex < 2:
  854. return self.mImgList[:5]
  855. elif currIndex > len(self.mImgList) - 3:
  856. return self.mImgList[-5:]
  857. else:
  858. return self.mImgList[currIndex - 2: currIndex + 3]
  859. # Tzutalin 20160906 : Add file list and dock to move faster
  860. def fileitemDoubleClicked(self, item=None):
  861. self.currIndex = self.mImgList.index(ustr(os.path.join(os.path.abspath(self.dirname), item.text())))
  862. filename = self.mImgList[self.currIndex]
  863. if filename:
  864. self.mImgList5 = self.indexTo5Files(self.currIndex)
  865. # self.additems5(None)
  866. self.loadFile(filename)
  867. def iconitemDoubleClicked(self, item=None):
  868. self.currIndex = self.mImgList.index(ustr(os.path.join(item.toolTip())))
  869. filename = self.mImgList[self.currIndex]
  870. if filename:
  871. self.mImgList5 = self.indexTo5Files(self.currIndex)
  872. # self.additems5(None)
  873. self.loadFile(filename)
  874. def CanvasSizeChange(self):
  875. if len(self.mImgList) > 0 and self.imageSlider.hasFocus():
  876. self.zoomWidget.setValue(self.imageSlider.value())
  877. def shapeSelectionChanged(self, selected_shapes):
  878. self._noSelectionSlot = True
  879. for shape in self.canvas.selectedShapes:
  880. shape.selected = False
  881. self.labelList.clearSelection()
  882. self.indexList.clearSelection()
  883. self.canvas.selectedShapes = selected_shapes
  884. for shape in self.canvas.selectedShapes:
  885. shape.selected = True
  886. self.shapesToItems[shape].setSelected(True)
  887. self.shapesToItemsbox[shape].setSelected(True)
  888. index = self.labelList.indexFromItem(self.shapesToItems[shape]).row()
  889. self.indexList.item(index).setSelected(True)
  890. self.labelList.scrollToItem(self.currentItem()) # QAbstractItemView.EnsureVisible
  891. # map current label item to index item and select it
  892. index = self.labelList.indexFromItem(self.currentItem()).row()
  893. self.indexList.scrollToItem(self.indexList.item(index))
  894. self.BoxList.scrollToItem(self.currentBox())
  895. if self.kie_mode:
  896. if len(self.canvas.selectedShapes) == 1 and self.keyList.count() > 0:
  897. selected_key_item_row = self.keyList.findItemsByLabel(self.canvas.selectedShapes[0].key_cls,
  898. get_row=True)
  899. if isinstance(selected_key_item_row, list) and len(selected_key_item_row) == 0:
  900. key_text = self.canvas.selectedShapes[0].key_cls
  901. item = self.keyList.createItemFromLabel(key_text)
  902. self.keyList.addItem(item)
  903. rgb = self._get_rgb_by_label(key_text, self.kie_mode)
  904. self.keyList.setItemLabel(item, key_text, rgb)
  905. selected_key_item_row = self.keyList.findItemsByLabel(self.canvas.selectedShapes[0].key_cls,
  906. get_row=True)
  907. self.keyList.setCurrentRow(selected_key_item_row)
  908. self._noSelectionSlot = False
  909. n_selected = len(selected_shapes)
  910. self.actions.singleRere.setEnabled(n_selected)
  911. self.actions.cellreRec.setEnabled(n_selected)
  912. self.actions.delete.setEnabled(n_selected)
  913. self.actions.copy.setEnabled(n_selected)
  914. self.actions.edit.setEnabled(n_selected == 1)
  915. self.actions.lock.setEnabled(n_selected)
  916. self.actions.change_cls.setEnabled(n_selected)
  917. def addLabel(self, shape):
  918. shape.paintLabel = self.displayLabelOption.isChecked()
  919. shape.paintIdx = self.displayIndexOption.isChecked()
  920. item = HashableQListWidgetItem(shape.label)
  921. # current difficult checkbox is disenble
  922. # item.setFlags(item.flags() | Qt.ItemIsUserCheckable)
  923. # item.setCheckState(Qt.Unchecked) if shape.difficult else item.setCheckState(Qt.Checked)
  924. # Checked means difficult is False
  925. # item.setBackground(generateColorByText(shape.label))
  926. self.itemsToShapes[item] = shape
  927. self.shapesToItems[shape] = item
  928. # add current label item index before label string
  929. current_index = QListWidgetItem(str(self.labelList.count()))
  930. current_index.setTextAlignment(Qt.AlignHCenter)
  931. self.indexList.addItem(current_index)
  932. self.labelList.addItem(item)
  933. # print('item in add label is ',[(p.x(), p.y()) for p in shape.points], shape.label)
  934. # ADD for box
  935. item = HashableQListWidgetItem(str([(int(p.x()), int(p.y())) for p in shape.points]))
  936. self.itemsToShapesbox[item] = shape
  937. self.shapesToItemsbox[shape] = item
  938. self.BoxList.addItem(item)
  939. for action in self.actions.onShapesPresent:
  940. action.setEnabled(True)
  941. self.updateComboBox()
  942. # update show counting
  943. self.BoxListDock.setWindowTitle(self.BoxListDockName + f" ({self.BoxList.count()})")
  944. self.labelListDock.setWindowTitle(self.labelListDockName + f" ({self.labelList.count()})")
  945. def remLabels(self, shapes):
  946. if shapes is None:
  947. # print('rm empty label')
  948. return
  949. for shape in shapes:
  950. item = self.shapesToItems[shape]
  951. self.labelList.takeItem(self.labelList.row(item))
  952. del self.shapesToItems[shape]
  953. del self.itemsToShapes[item]
  954. self.updateComboBox()
  955. # ADD:
  956. item = self.shapesToItemsbox[shape]
  957. self.BoxList.takeItem(self.BoxList.row(item))
  958. del self.shapesToItemsbox[shape]
  959. del self.itemsToShapesbox[item]
  960. self.updateComboBox()
  961. self.updateIndexList()
  962. def loadLabels(self, shapes):
  963. s = []
  964. shape_index = 0
  965. for label, points, line_color, key_cls, difficult in shapes:
  966. shape = Shape(label=label, line_color=line_color, key_cls=key_cls)
  967. for x, y in points:
  968. # Ensure the labels are within the bounds of the image. If not, fix them.
  969. x, y, snapped = self.canvas.snapPointToCanvas(x, y)
  970. if snapped:
  971. self.setDirty()
  972. shape.addPoint(QPointF(x, y))
  973. shape.difficult = difficult
  974. shape.idx = shape_index
  975. shape_index += 1
  976. # shape.locked = False
  977. shape.close()
  978. s.append(shape)
  979. self._update_shape_color(shape)
  980. self.addLabel(shape)
  981. self.updateComboBox()
  982. self.canvas.loadShapes(s)
  983. def singleLabel(self, shape):
  984. if shape is None:
  985. # print('rm empty label')
  986. return
  987. item = self.shapesToItems[shape]
  988. item.setText(shape.label)
  989. self.updateComboBox()
  990. # ADD:
  991. item = self.shapesToItemsbox[shape]
  992. item.setText(str([(int(p.x()), int(p.y())) for p in shape.points]))
  993. self.updateComboBox()
  994. def updateComboBox(self):
  995. # Get the unique labels and add them to the Combobox.
  996. itemsTextList = [str(self.labelList.item(i).text()) for i in range(self.labelList.count())]
  997. uniqueTextList = list(set(itemsTextList))
  998. # Add a null row for showing all the labels
  999. uniqueTextList.append("")
  1000. uniqueTextList.sort()
  1001. # self.comboBox.update_items(uniqueTextList)
  1002. def updateIndexList(self):
  1003. self.indexList.clear()
  1004. for i in range(self.labelList.count()):
  1005. string = QListWidgetItem(str(i))
  1006. string.setTextAlignment(Qt.AlignHCenter)
  1007. self.indexList.addItem(string)
  1008. def saveLabels(self, annotationFilePath, mode='Auto'):
  1009. # Mode is Auto means that labels will be loaded from self.result_dic totally, which is the output of ocr model
  1010. annotationFilePath = ustr(annotationFilePath)
  1011. def format_shape(s):
  1012. # print('s in saveLabels is ',s)
  1013. return dict(label=s.label, # str
  1014. line_color=s.line_color.getRgb(),
  1015. fill_color=s.fill_color.getRgb(),
  1016. points=[(int(p.x()), int(p.y())) for p in s.points], # QPonitF
  1017. difficult=s.difficult,
  1018. key_cls=s.key_cls) # bool
  1019. if mode == 'Auto':
  1020. shapes = []
  1021. else:
  1022. shapes = [format_shape(shape) for shape in self.canvas.shapes if shape.line_color != DEFAULT_LOCK_COLOR]
  1023. # Can add differrent annotation formats here
  1024. for box in self.result_dic:
  1025. trans_dic = {"label": box[1][0], "points": box[0], "difficult": False}
  1026. if self.kie_mode:
  1027. if len(box) == 3:
  1028. trans_dic.update({"key_cls": box[2]})
  1029. else:
  1030. trans_dic.update({"key_cls": "None"})
  1031. if trans_dic["label"] == "" and mode == 'Auto':
  1032. continue
  1033. shapes.append(trans_dic)
  1034. try:
  1035. trans_dic = []
  1036. for box in shapes:
  1037. trans_dict = {"transcription": box['label'], "points": box['points'], "difficult": box['difficult']}
  1038. if self.kie_mode:
  1039. trans_dict.update({"key_cls": box['key_cls']})
  1040. trans_dic.append(trans_dict)
  1041. self.PPlabel[annotationFilePath] = trans_dic
  1042. if mode == 'Auto':
  1043. self.Cachelabel[annotationFilePath] = trans_dic
  1044. # else:
  1045. # self.labelFile.save(annotationFilePath, shapes, self.filePath, self.imageData,
  1046. # self.lineColor.getRgb(), self.fillColor.getRgb())
  1047. # print('Image:{0} -> Annotation:{1}'.format(self.filePath, annotationFilePath))
  1048. return True
  1049. except:
  1050. self.errorMessage(u'Error saving label data', u'Error saving label data')
  1051. return False
  1052. def copySelectedShape(self):
  1053. for shape in self.canvas.copySelectedShape():
  1054. self.addLabel(shape)
  1055. # fix copy and delete
  1056. # self.shapeSelectionChanged(True)
  1057. def move_scrollbar(self, value):
  1058. self.labelListBar.setValue(value)
  1059. self.indexListBar.setValue(value)
  1060. def labelSelectionChanged(self):
  1061. if self._noSelectionSlot:
  1062. return
  1063. if self.canvas.editing():
  1064. selected_shapes = []
  1065. for item in self.labelList.selectedItems():
  1066. selected_shapes.append(self.itemsToShapes[item])
  1067. if selected_shapes:
  1068. self.canvas.selectShapes(selected_shapes)
  1069. else:
  1070. self.canvas.deSelectShape()
  1071. def indexSelectionChanged(self):
  1072. if self._noSelectionSlot:
  1073. return
  1074. if self.canvas.editing():
  1075. selected_shapes = []
  1076. for item in self.indexList.selectedItems():
  1077. # map index item to label item
  1078. index = self.indexList.indexFromItem(item).row()
  1079. item = self.labelList.item(index)
  1080. selected_shapes.append(self.itemsToShapes[item])
  1081. if selected_shapes:
  1082. self.canvas.selectShapes(selected_shapes)
  1083. else:
  1084. self.canvas.deSelectShape()
  1085. def boxSelectionChanged(self):
  1086. if self._noSelectionSlot:
  1087. # self.BoxList.scrollToItem(self.currentBox(), QAbstractItemView.PositionAtCenter)
  1088. return
  1089. if self.canvas.editing():
  1090. selected_shapes = []
  1091. for item in self.BoxList.selectedItems():
  1092. selected_shapes.append(self.itemsToShapesbox[item])
  1093. if selected_shapes:
  1094. self.canvas.selectShapes(selected_shapes)
  1095. else:
  1096. self.canvas.deSelectShape()
  1097. def labelItemChanged(self, item):
  1098. # avoid accidentally triggering the itemChanged siganl with unhashable item
  1099. # Unknown trigger condition
  1100. if type(item) == HashableQListWidgetItem:
  1101. shape = self.itemsToShapes[item]
  1102. label = item.text()
  1103. if label != shape.label:
  1104. shape.label = item.text()
  1105. # shape.line_color = generateColorByText(shape.label)
  1106. self.setDirty()
  1107. elif not ((item.checkState() == Qt.Unchecked) ^ (not shape.difficult)):
  1108. shape.difficult = True if item.checkState() == Qt.Unchecked else False
  1109. self.setDirty()
  1110. else: # User probably changed item visibility
  1111. self.canvas.setShapeVisible(shape, True) # item.checkState() == Qt.Checked
  1112. # self.actions.save.setEnabled(True)
  1113. else:
  1114. print('enter labelItemChanged slot with unhashable item: ', item, item.text())
  1115. def drag_drop_happened(self):
  1116. '''
  1117. label list drag drop signal slot
  1118. '''
  1119. # print('___________________drag_drop_happened_______________')
  1120. # should only select single item
  1121. for item in self.labelList.selectedItems():
  1122. newIndex = self.labelList.indexFromItem(item).row()
  1123. # only support drag_drop one item
  1124. assert len(self.canvas.selectedShapes) > 0
  1125. for shape in self.canvas.selectedShapes:
  1126. selectedShapeIndex = shape.idx
  1127. if newIndex == selectedShapeIndex:
  1128. return
  1129. # move corresponding item in shape list
  1130. shape = self.canvas.shapes.pop(selectedShapeIndex)
  1131. self.canvas.shapes.insert(newIndex, shape)
  1132. # update bbox index
  1133. self.canvas.updateShapeIndex()
  1134. # boxList update simultaneously
  1135. item = self.BoxList.takeItem(selectedShapeIndex)
  1136. self.BoxList.insertItem(newIndex, item)
  1137. # changes happen
  1138. self.setDirty()
  1139. # Callback functions:
  1140. def newShape(self, value=True):
  1141. """Pop-up and give focus to the label editor.
  1142. position MUST be in global coordinates.
  1143. """
  1144. if len(self.labelHist) > 0:
  1145. self.labelDialog = LabelDialog(parent=self, listItem=self.labelHist)
  1146. if value:
  1147. text = self.labelDialog.popUp(text=self.prevLabelText)
  1148. self.lastLabel = text
  1149. else:
  1150. text = self.prevLabelText
  1151. if text is not None:
  1152. self.prevLabelText = self.stringBundle.getString('tempLabel')
  1153. shape = self.canvas.setLastLabel(text, None, None, None) # generate_color, generate_color
  1154. if self.kie_mode:
  1155. key_text, _ = self.keyDialog.popUp(self.key_previous_text)
  1156. if key_text is not None:
  1157. shape = self.canvas.setLastLabel(text, None, None, key_text) # generate_color, generate_color
  1158. self.key_previous_text = key_text
  1159. if not self.keyList.findItemsByLabel(key_text):
  1160. item = self.keyList.createItemFromLabel(key_text)
  1161. self.keyList.addItem(item)
  1162. rgb = self._get_rgb_by_label(key_text, self.kie_mode)
  1163. self.keyList.setItemLabel(item, key_text, rgb)
  1164. self._update_shape_color(shape)
  1165. self.keyDialog.addLabelHistory(key_text)
  1166. self.addLabel(shape)
  1167. if self.beginner(): # Switch to edit mode.
  1168. self.canvas.setEditing(True)
  1169. self.actions.create.setEnabled(True)
  1170. self.actions.createpoly.setEnabled(True)
  1171. self.actions.undoLastPoint.setEnabled(False)
  1172. self.actions.undo.setEnabled(True)
  1173. else:
  1174. self.actions.editMode.setEnabled(True)
  1175. self.setDirty()
  1176. else:
  1177. # self.canvas.undoLastLine()
  1178. self.canvas.resetAllLines()
  1179. def _update_shape_color(self, shape):
  1180. r, g, b = self._get_rgb_by_label(shape.key_cls, self.kie_mode)
  1181. shape.line_color = QColor(r, g, b)
  1182. shape.vertex_fill_color = QColor(r, g, b)
  1183. shape.hvertex_fill_color = QColor(255, 255, 255)
  1184. shape.fill_color = QColor(r, g, b, 128)
  1185. shape.select_line_color = QColor(255, 255, 255)
  1186. shape.select_fill_color = QColor(r, g, b, 155)
  1187. def _get_rgb_by_label(self, label, kie_mode):
  1188. shift_auto_shape_color = 2 # use for random color
  1189. if kie_mode and label != "None":
  1190. item = self.keyList.findItemsByLabel(label)[0]
  1191. label_id = self.keyList.indexFromItem(item).row() + 1
  1192. label_id += shift_auto_shape_color
  1193. return LABEL_COLORMAP[label_id % len(LABEL_COLORMAP)]
  1194. else:
  1195. return (0, 255, 0)
  1196. def scrollRequest(self, delta, orientation):
  1197. units = - delta / (8 * 15)
  1198. bar = self.scrollBars[orientation]
  1199. bar.setValue(bar.value() + bar.singleStep() * units)
  1200. def setZoom(self, value):
  1201. self.actions.fitWidth.setChecked(False)
  1202. self.actions.fitWindow.setChecked(False)
  1203. self.zoomMode = self.MANUAL_ZOOM
  1204. self.zoomWidget.setValue(value)
  1205. def addZoom(self, increment=10):
  1206. self.setZoom(self.zoomWidget.value() + increment)
  1207. self.imageSlider.setValue(self.zoomWidget.value() + increment) # set zoom slider value
  1208. def zoomRequest(self, delta):
  1209. # get the current scrollbar positions
  1210. # calculate the percentages ~ coordinates
  1211. h_bar = self.scrollBars[Qt.Horizontal]
  1212. v_bar = self.scrollBars[Qt.Vertical]
  1213. # get the current maximum, to know the difference after zooming
  1214. h_bar_max = h_bar.maximum()
  1215. v_bar_max = v_bar.maximum()
  1216. # get the cursor position and canvas size
  1217. # calculate the desired movement from 0 to 1
  1218. # where 0 = move left
  1219. # 1 = move right
  1220. # up and down analogous
  1221. cursor = QCursor()
  1222. pos = cursor.pos()
  1223. relative_pos = QWidget.mapFromGlobal(self, pos)
  1224. cursor_x = relative_pos.x()
  1225. cursor_y = relative_pos.y()
  1226. w = self.scrollArea.width()
  1227. h = self.scrollArea.height()
  1228. # the scaling from 0 to 1 has some padding
  1229. # you don't have to hit the very leftmost pixel for a maximum-left movement
  1230. margin = 0.1
  1231. move_x = (cursor_x - margin * w) / (w - 2 * margin * w)
  1232. move_y = (cursor_y - margin * h) / (h - 2 * margin * h)
  1233. # clamp the values from 0 to 1
  1234. move_x = min(max(move_x, 0), 1)
  1235. move_y = min(max(move_y, 0), 1)
  1236. # zoom in
  1237. units = delta / (8 * 15)
  1238. scale = 10
  1239. self.addZoom(scale * units)
  1240. # get the difference in scrollbar values
  1241. # this is how far we can move
  1242. d_h_bar_max = h_bar.maximum() - h_bar_max
  1243. d_v_bar_max = v_bar.maximum() - v_bar_max
  1244. # get the new scrollbar values
  1245. new_h_bar_value = h_bar.value() + move_x * d_h_bar_max
  1246. new_v_bar_value = v_bar.value() + move_y * d_v_bar_max
  1247. h_bar.setValue(new_h_bar_value)
  1248. v_bar.setValue(new_v_bar_value)
  1249. def setFitWindow(self, value=True):
  1250. if value:
  1251. self.actions.fitWidth.setChecked(False)
  1252. self.zoomMode = self.FIT_WINDOW if value else self.MANUAL_ZOOM
  1253. self.adjustScale()
  1254. def setFitWidth(self, value=True):
  1255. if value:
  1256. self.actions.fitWindow.setChecked(False)
  1257. self.zoomMode = self.FIT_WIDTH if value else self.MANUAL_ZOOM
  1258. self.adjustScale()
  1259. def togglePolygons(self, value):
  1260. for item, shape in self.itemsToShapes.items():
  1261. self.canvas.setShapeVisible(shape, value)
  1262. def loadFile(self, filePath=None):
  1263. """Load the specified file, or the last opened file if None."""
  1264. if self.dirty:
  1265. self.mayContinue()
  1266. self.resetState()
  1267. self.canvas.setEnabled(False)
  1268. if filePath is None:
  1269. filePath = self.settings.get(SETTING_FILENAME)
  1270. # Make sure that filePath is a regular python string, rather than QString
  1271. filePath = ustr(filePath)
  1272. # Fix bug: An index error after select a directory when open a new file.
  1273. unicodeFilePath = ustr(filePath)
  1274. # unicodeFilePath = os.path.abspath(unicodeFilePath)
  1275. # Tzutalin 20160906 : Add file list and dock to move faster
  1276. # Highlight the file item
  1277. if unicodeFilePath and self.fileListWidget.count() > 0:
  1278. if unicodeFilePath in self.mImgList:
  1279. index = self.mImgList.index(unicodeFilePath)
  1280. fileWidgetItem = self.fileListWidget.item(index)
  1281. print('unicodeFilePath is', unicodeFilePath)
  1282. fileWidgetItem.setSelected(True)
  1283. self.iconlist.clear()
  1284. self.additems5(None)
  1285. for i in range(5):
  1286. item_tooltip = self.iconlist.item(i).toolTip()
  1287. # print(i,"---",item_tooltip)
  1288. if item_tooltip == ustr(filePath):
  1289. titem = self.iconlist.item(i)
  1290. titem.setSelected(True)
  1291. self.iconlist.scrollToItem(titem)
  1292. break
  1293. else:
  1294. self.fileListWidget.clear()
  1295. self.mImgList.clear()
  1296. self.iconlist.clear()
  1297. # if unicodeFilePath and self.iconList.count() > 0:
  1298. # if unicodeFilePath in self.mImgList:
  1299. if unicodeFilePath and os.path.exists(unicodeFilePath):
  1300. self.canvas.verified = False
  1301. cvimg = cv2.imdecode(np.fromfile(unicodeFilePath, dtype=np.uint8), 1)
  1302. height, width, depth = cvimg.shape
  1303. cvimg = cv2.cvtColor(cvimg, cv2.COLOR_BGR2RGB)
  1304. image = QImage(cvimg.data, width, height, width * depth, QImage.Format_RGB888)
  1305. if image.isNull():
  1306. self.errorMessage(u'Error opening file',
  1307. u"<p>Make sure <i>%s</i> is a valid image file." % unicodeFilePath)
  1308. self.status("Error reading %s" % unicodeFilePath)
  1309. return False
  1310. self.status("Loaded %s" % os.path.basename(unicodeFilePath))
  1311. self.image = image
  1312. self.filePath = unicodeFilePath
  1313. self.canvas.loadPixmap(QPixmap.fromImage(image))
  1314. if self.validFilestate(filePath) is True:
  1315. self.setClean()
  1316. else:
  1317. self.dirty = False
  1318. self.actions.save.setEnabled(True)
  1319. if len(self.canvas.lockedShapes) != 0:
  1320. self.actions.save.setEnabled(True)
  1321. self.setDirty()
  1322. self.canvas.setEnabled(True)
  1323. self.adjustScale(initial=True)
  1324. self.paintCanvas()
  1325. self.addRecentFile(self.filePath)
  1326. self.toggleActions(True)
  1327. self.showBoundingBoxFromPPlabel(filePath)
  1328. self.setWindowTitle(__appname__ + ' ' + filePath)
  1329. # Default : select last item if there is at least one item
  1330. if self.labelList.count():
  1331. self.labelList.setCurrentItem(self.labelList.item(self.labelList.count() - 1))
  1332. self.labelList.item(self.labelList.count() - 1).setSelected(True)
  1333. self.indexList.item(self.labelList.count() - 1).setSelected(True)
  1334. # show file list image count
  1335. select_indexes = self.fileListWidget.selectedIndexes()
  1336. if len(select_indexes) > 0:
  1337. self.fileDock.setWindowTitle(self.fileListName + f" ({select_indexes[0].row() + 1}"
  1338. f"/{self.fileListWidget.count()})")
  1339. # update show counting
  1340. self.BoxListDock.setWindowTitle(self.BoxListDockName + f" ({self.BoxList.count()})")
  1341. self.labelListDock.setWindowTitle(self.labelListDockName + f" ({self.labelList.count()})")
  1342. self.canvas.setFocus(True)
  1343. return True
  1344. return False
  1345. def showBoundingBoxFromPPlabel(self, filePath):
  1346. width, height = self.image.width(), self.image.height()
  1347. imgidx = self.getImglabelidx(filePath)
  1348. shapes = []
  1349. # box['ratio'] of the shapes saved in lockedShapes contains the ratio of the
  1350. # four corner coordinates of the shapes to the height and width of the image
  1351. for box in self.canvas.lockedShapes:
  1352. key_cls = 'None' if not self.kie_mode else box['key_cls']
  1353. if self.canvas.isInTheSameImage:
  1354. shapes.append((box['transcription'], [[s[0] * width, s[1] * height] for s in box['ratio']],
  1355. DEFAULT_LOCK_COLOR, key_cls, box['difficult']))
  1356. else:
  1357. shapes.append(('锁定框:待检测', [[s[0] * width, s[1] * height] for s in box['ratio']],
  1358. DEFAULT_LOCK_COLOR, key_cls, box['difficult']))
  1359. if imgidx in self.PPlabel.keys():
  1360. for box in self.PPlabel[imgidx]:
  1361. key_cls = 'None' if not self.kie_mode else box.get('key_cls', 'None')
  1362. shapes.append((box['transcription'], box['points'], None, key_cls, box.get('difficult', False)))
  1363. if shapes != []:
  1364. self.loadLabels(shapes)
  1365. self.canvas.verified = False
  1366. def validFilestate(self, filePath):
  1367. if filePath not in self.fileStatedict.keys():
  1368. return None
  1369. elif self.fileStatedict[filePath] == 1:
  1370. return True
  1371. else:
  1372. return False
  1373. def resizeEvent(self, event):
  1374. if self.canvas and not self.image.isNull() \
  1375. and self.zoomMode != self.MANUAL_ZOOM:
  1376. self.adjustScale()
  1377. super(MainWindow, self).resizeEvent(event)
  1378. def paintCanvas(self):
  1379. assert not self.image.isNull(), "cannot paint null image"
  1380. self.canvas.scale = 0.01 * self.zoomWidget.value()
  1381. self.canvas.adjustSize()
  1382. self.canvas.update()
  1383. def adjustScale(self, initial=False):
  1384. value = self.scalers[self.FIT_WINDOW if initial else self.zoomMode]()
  1385. self.zoomWidget.setValue(int(100 * value))
  1386. self.imageSlider.setValue(self.zoomWidget.value()) # set zoom slider value
  1387. def scaleFitWindow(self):
  1388. """Figure out the size of the pixmap in order to fit the main widget."""
  1389. e = 2.0 # So that no scrollbars are generated.
  1390. w1 = self.centralWidget().width() - e
  1391. h1 = self.centralWidget().height() - e - 110
  1392. a1 = w1 / h1
  1393. # Calculate a new scale value based on the pixmap's aspect ratio.
  1394. w2 = self.canvas.pixmap.width() - 0.0
  1395. h2 = self.canvas.pixmap.height() - 0.0
  1396. a2 = w2 / h2
  1397. return w1 / w2 if a2 >= a1 else h1 / h2
  1398. def scaleFitWidth(self):
  1399. # The epsilon does not seem to work too well here.
  1400. w = self.centralWidget().width() - 2.0
  1401. return w / self.canvas.pixmap.width()
  1402. def closeEvent(self, event):
  1403. if not self.mayContinue():
  1404. event.ignore()
  1405. else:
  1406. settings = self.settings
  1407. # If it loads images from dir, don't load it at the beginning
  1408. if self.dirname is None:
  1409. settings[SETTING_FILENAME] = self.filePath if self.filePath else ''
  1410. else:
  1411. settings[SETTING_FILENAME] = ''
  1412. settings[SETTING_WIN_SIZE] = self.size()
  1413. settings[SETTING_WIN_POSE] = self.pos()
  1414. settings[SETTING_WIN_STATE] = self.saveState()
  1415. settings[SETTING_LINE_COLOR] = self.lineColor
  1416. settings[SETTING_FILL_COLOR] = self.fillColor
  1417. settings[SETTING_RECENT_FILES] = self.recentFiles
  1418. settings[SETTING_ADVANCE_MODE] = not self._beginner
  1419. if self.defaultSaveDir and os.path.exists(self.defaultSaveDir):
  1420. settings[SETTING_SAVE_DIR] = ustr(self.defaultSaveDir)
  1421. else:
  1422. settings[SETTING_SAVE_DIR] = ''
  1423. if self.lastOpenDir and os.path.exists(self.lastOpenDir):
  1424. settings[SETTING_LAST_OPEN_DIR] = self.lastOpenDir
  1425. else:
  1426. settings[SETTING_LAST_OPEN_DIR] = ''
  1427. settings[SETTING_PAINT_LABEL] = self.displayLabelOption.isChecked()
  1428. settings[SETTING_PAINT_INDEX] = self.displayIndexOption.isChecked()
  1429. settings[SETTING_DRAW_SQUARE] = self.drawSquaresOption.isChecked()
  1430. settings.save()
  1431. try:
  1432. self.saveLabelFile()
  1433. except:
  1434. pass
  1435. def loadRecent(self, filename):
  1436. if self.mayContinue():
  1437. print(filename, "======")
  1438. self.loadFile(filename)
  1439. def scanAllImages(self, folderPath):
  1440. extensions = ['.%s' % fmt.data().decode("ascii").lower() for fmt in QImageReader.supportedImageFormats()]
  1441. images = []
  1442. for file in os.listdir(folderPath):
  1443. if file.lower().endswith(tuple(extensions)):
  1444. relativePath = os.path.join(folderPath, file)
  1445. path = ustr(os.path.abspath(relativePath))
  1446. images.append(path)
  1447. natural_sort(images, key=lambda x: x.lower())
  1448. return images
  1449. def openDirDialog(self, _value=False, dirpath=None, silent=False):
  1450. if not self.mayContinue():
  1451. return
  1452. defaultOpenDirPath = dirpath if dirpath else '.'
  1453. if self.lastOpenDir and os.path.exists(self.lastOpenDir):
  1454. defaultOpenDirPath = self.lastOpenDir
  1455. else:
  1456. defaultOpenDirPath = os.path.dirname(self.filePath) if self.filePath else '.'
  1457. if silent != True:
  1458. targetDirPath = ustr(QFileDialog.getExistingDirectory(self,
  1459. '%s - Open Directory' % __appname__,
  1460. defaultOpenDirPath,
  1461. QFileDialog.ShowDirsOnly | QFileDialog.DontResolveSymlinks))
  1462. else:
  1463. targetDirPath = ustr(defaultOpenDirPath)
  1464. self.lastOpenDir = targetDirPath
  1465. self.importDirImages(targetDirPath)
  1466. def openDatasetDirDialog(self):
  1467. if self.lastOpenDir and os.path.exists(self.lastOpenDir):
  1468. if platform.system() == 'Windows':
  1469. os.startfile(self.lastOpenDir)
  1470. else:
  1471. os.system('open ' + os.path.normpath(self.lastOpenDir))
  1472. defaultOpenDirPath = self.lastOpenDir
  1473. else:
  1474. if self.lang == 'ch':
  1475. self.msgBox.warning(self, "提示", "\n 原文件夹已不存在,请从新选择数据集路径!")
  1476. else:
  1477. self.msgBox.warning(self, "Warn",
  1478. "\n The original folder no longer exists, please choose the data set path again!")
  1479. self.actions.open_dataset_dir.setEnabled(False)
  1480. defaultOpenDirPath = os.path.dirname(self.filePath) if self.filePath else '.'
  1481. def init_key_list(self, label_dict):
  1482. if not self.kie_mode:
  1483. return
  1484. # load key_cls
  1485. for image, info in label_dict.items():
  1486. for box in info:
  1487. if "key_cls" not in box:
  1488. box.update({"key_cls": "None"})
  1489. self.existed_key_cls_set.add(box["key_cls"])
  1490. if len(self.existed_key_cls_set) > 0:
  1491. for key_text in self.existed_key_cls_set:
  1492. if not self.keyList.findItemsByLabel(key_text):
  1493. item = self.keyList.createItemFromLabel(key_text)
  1494. self.keyList.addItem(item)
  1495. rgb = self._get_rgb_by_label(key_text, self.kie_mode)
  1496. self.keyList.setItemLabel(item, key_text, rgb)
  1497. if self.keyDialog is None:
  1498. # key list dialog
  1499. self.keyDialog = KeyDialog(
  1500. text=self.key_dialog_tip,
  1501. parent=self,
  1502. labels=self.existed_key_cls_set,
  1503. sort_labels=True,
  1504. show_text_field=True,
  1505. completion="startswith",
  1506. fit_to_content={'column': True, 'row': False},
  1507. flags=None
  1508. )
  1509. def importDirImages(self, dirpath, isDelete=False):
  1510. if not self.mayContinue() or not dirpath:
  1511. return
  1512. if self.defaultSaveDir and self.defaultSaveDir != dirpath:
  1513. self.saveLabelFile()
  1514. if not isDelete:
  1515. self.loadFilestate(dirpath)
  1516. self.PPlabelpath = dirpath + '/Label.txt'
  1517. self.PPlabel = self.loadLabelFile(self.PPlabelpath)
  1518. self.Cachelabelpath = dirpath + '/Cache.cach'
  1519. self.Cachelabel = self.loadLabelFile(self.Cachelabelpath)
  1520. if self.Cachelabel:
  1521. self.PPlabel = dict(self.Cachelabel, **self.PPlabel)
  1522. self.init_key_list(self.PPlabel)
  1523. self.lastOpenDir = dirpath
  1524. self.dirname = dirpath
  1525. self.defaultSaveDir = dirpath
  1526. self.statusBar().showMessage('%s started. Annotation will be saved to %s' %
  1527. (__appname__, self.defaultSaveDir))
  1528. self.statusBar().show()
  1529. self.filePath = None
  1530. self.fileListWidget.clear()
  1531. self.mImgList = self.scanAllImages(dirpath)
  1532. self.mImgList5 = self.mImgList[:5]
  1533. self.openNextImg()
  1534. doneicon = newIcon('done')
  1535. closeicon = newIcon('close')
  1536. for imgPath in self.mImgList:
  1537. filename = os.path.basename(imgPath)
  1538. if self.validFilestate(imgPath) is True:
  1539. item = QListWidgetItem(doneicon, filename)
  1540. else:
  1541. item = QListWidgetItem(closeicon, filename)
  1542. self.fileListWidget.addItem(item)
  1543. print('DirPath in importDirImages is', dirpath)
  1544. self.iconlist.clear()
  1545. self.additems5(dirpath)
  1546. self.changeFileFolder = True
  1547. self.haveAutoReced = False
  1548. self.AutoRecognition.setEnabled(True)
  1549. self.reRecogButton.setEnabled(True)
  1550. self.tableRecButton.setEnabled(True)
  1551. self.actions.AutoRec.setEnabled(True)
  1552. self.actions.reRec.setEnabled(True)
  1553. self.actions.tableRec.setEnabled(True)
  1554. self.actions.open_dataset_dir.setEnabled(True)
  1555. self.actions.rotateLeft.setEnabled(True)
  1556. self.actions.rotateRight.setEnabled(True)
  1557. self.fileListWidget.setCurrentRow(0) # set list index to first
  1558. self.fileDock.setWindowTitle(self.fileListName + f" (1/{self.fileListWidget.count()})") # show image count
  1559. def openPrevImg(self, _value=False):
  1560. if len(self.mImgList) <= 0:
  1561. return
  1562. if self.filePath is None:
  1563. return
  1564. currIndex = self.mImgList.index(self.filePath)
  1565. self.mImgList5 = self.mImgList[:5]
  1566. if currIndex - 1 >= 0:
  1567. filename = self.mImgList[currIndex - 1]
  1568. self.mImgList5 = self.indexTo5Files(currIndex - 1)
  1569. if filename:
  1570. self.loadFile(filename)
  1571. def openNextImg(self, _value=False):
  1572. if not self.mayContinue():
  1573. return
  1574. if len(self.mImgList) <= 0:
  1575. return
  1576. filename = None
  1577. if self.filePath is None:
  1578. filename = self.mImgList[0]
  1579. self.mImgList5 = self.mImgList[:5]
  1580. else:
  1581. currIndex = self.mImgList.index(self.filePath)
  1582. if currIndex + 1 < len(self.mImgList):
  1583. filename = self.mImgList[currIndex + 1]
  1584. self.mImgList5 = self.indexTo5Files(currIndex + 1)
  1585. else:
  1586. self.mImgList5 = self.indexTo5Files(currIndex)
  1587. if filename:
  1588. print('file name in openNext is ', filename)
  1589. self.loadFile(filename)
  1590. def updateFileListIcon(self, filename):
  1591. pass
  1592. def saveFile(self, _value=False, mode='Manual'):
  1593. # Manual mode is used for users click "Save" manually,which will change the state of the image
  1594. if self.filePath:
  1595. imgidx = self.getImglabelidx(self.filePath)
  1596. self._saveFile(imgidx, mode=mode)
  1597. def saveLockedShapes(self):
  1598. self.canvas.lockedShapes = []
  1599. self.canvas.selectedShapes = []
  1600. for s in self.canvas.shapes:
  1601. if s.line_color == DEFAULT_LOCK_COLOR:
  1602. self.canvas.selectedShapes.append(s)
  1603. self.lockSelectedShape()
  1604. for s in self.canvas.shapes:
  1605. if s.line_color == DEFAULT_LOCK_COLOR:
  1606. self.canvas.selectedShapes.remove(s)
  1607. self.canvas.shapes.remove(s)
  1608. def _saveFile(self, annotationFilePath, mode='Manual'):
  1609. if len(self.canvas.lockedShapes) != 0:
  1610. self.saveLockedShapes()
  1611. if mode == 'Manual':
  1612. self.result_dic_locked = []
  1613. img = cv2.imread(self.filePath)
  1614. width, height = self.image.width(), self.image.height()
  1615. for shape in self.canvas.lockedShapes:
  1616. box = [[int(p[0] * width), int(p[1] * height)] for p in shape['ratio']]
  1617. # assert len(box) == 4
  1618. result = [(shape['transcription'], 1)]
  1619. result.insert(0, box)
  1620. self.result_dic_locked.append(result)
  1621. self.result_dic += self.result_dic_locked
  1622. self.result_dic_locked = []
  1623. if annotationFilePath and self.saveLabels(annotationFilePath, mode=mode):
  1624. self.setClean()
  1625. self.statusBar().showMessage('Saved to %s' % annotationFilePath)
  1626. self.statusBar().show()
  1627. currIndex = self.mImgList.index(self.filePath)
  1628. item = self.fileListWidget.item(currIndex)
  1629. item.setIcon(newIcon('done'))
  1630. self.fileStatedict[self.filePath] = 1
  1631. if len(self.fileStatedict) % self.autoSaveNum == 0:
  1632. self.saveFilestate()
  1633. self.savePPlabel(mode='Auto')
  1634. self.fileListWidget.insertItem(int(currIndex), item)
  1635. if not self.canvas.isInTheSameImage:
  1636. self.openNextImg()
  1637. self.actions.saveRec.setEnabled(True)
  1638. self.actions.saveLabel.setEnabled(True)
  1639. self.actions.exportJSON.setEnabled(True)
  1640. elif mode == 'Auto':
  1641. if annotationFilePath and self.saveLabels(annotationFilePath, mode=mode):
  1642. self.setClean()
  1643. self.statusBar().showMessage('Saved to %s' % annotationFilePath)
  1644. self.statusBar().show()
  1645. def closeFile(self, _value=False):
  1646. if not self.mayContinue():
  1647. return
  1648. self.resetState()
  1649. self.setClean()
  1650. self.toggleActions(False)
  1651. self.canvas.setEnabled(False)
  1652. self.actions.saveAs.setEnabled(False)
  1653. def deleteImg(self):
  1654. deletePath = self.filePath
  1655. if deletePath is not None:
  1656. deleteInfo = self.deleteImgDialog()
  1657. if deleteInfo == QMessageBox.Yes:
  1658. if platform.system() == 'Windows':
  1659. from win32com.shell import shell, shellcon
  1660. shell.SHFileOperation((0, shellcon.FO_DELETE, deletePath, None,
  1661. shellcon.FOF_SILENT | shellcon.FOF_ALLOWUNDO | shellcon.FOF_NOCONFIRMATION,
  1662. None, None))
  1663. # linux
  1664. elif platform.system() == 'Linux':
  1665. cmd = 'trash ' + deletePath
  1666. os.system(cmd)
  1667. # macOS
  1668. elif platform.system() == 'Darwin':
  1669. import subprocess
  1670. absPath = os.path.abspath(deletePath).replace('\\', '\\\\').replace('"', '\\"')
  1671. cmd = ['osascript', '-e',
  1672. 'tell app "Finder" to move {the POSIX file "' + absPath + '"} to trash']
  1673. print(cmd)
  1674. subprocess.call(cmd, stdout=open(os.devnull, 'w'))
  1675. if self.filePath in self.fileStatedict.keys():
  1676. self.fileStatedict.pop(self.filePath)
  1677. imgidx = self.getImglabelidx(self.filePath)
  1678. if imgidx in self.PPlabel.keys():
  1679. self.PPlabel.pop(imgidx)
  1680. self.openNextImg()
  1681. self.importDirImages(self.lastOpenDir, isDelete=True)
  1682. def deleteImgDialog(self):
  1683. yes, cancel = QMessageBox.Yes, QMessageBox.Cancel
  1684. msg = u'The image will be deleted to the recycle bin'
  1685. return QMessageBox.warning(self, u'Attention', msg, yes | cancel)
  1686. def resetAll(self):
  1687. self.settings.reset()
  1688. self.close()
  1689. proc = QProcess()
  1690. proc.startDetached(os.path.abspath(__file__))
  1691. def mayContinue(self): #
  1692. if not self.dirty:
  1693. return True
  1694. else:
  1695. discardChanges = self.discardChangesDialog()
  1696. if discardChanges == QMessageBox.No:
  1697. return True
  1698. elif discardChanges == QMessageBox.Yes:
  1699. self.canvas.isInTheSameImage = True
  1700. self.saveFile()
  1701. self.canvas.isInTheSameImage = False
  1702. return True
  1703. else:
  1704. return False
  1705. def discardChangesDialog(self):
  1706. yes, no, cancel = QMessageBox.Yes, QMessageBox.No, QMessageBox.Cancel
  1707. if self.lang == 'ch':
  1708. msg = u'您有未保存的变更, 您想保存再继续吗?\n点击 "No" 丢弃所有未保存的变更.'
  1709. else:
  1710. msg = u'You have unsaved changes, would you like to save them and proceed?\nClick "No" to undo all changes.'
  1711. return QMessageBox.warning(self, u'Attention', msg, yes | no | cancel)
  1712. def errorMessage(self, title, message):
  1713. return QMessageBox.critical(self, title,
  1714. '<p><b>%s</b></p>%s' % (title, message))
  1715. def currentPath(self):
  1716. return os.path.dirname(self.filePath) if self.filePath else '.'
  1717. def chooseColor(self):
  1718. color = self.colorDialog.getColor(self.lineColor, u'Choose line color',
  1719. default=DEFAULT_LINE_COLOR)
  1720. if color:
  1721. self.lineColor = color
  1722. Shape.line_color = color
  1723. self.canvas.setDrawingColor(color)
  1724. self.canvas.update()
  1725. self.setDirty()
  1726. def deleteSelectedShape(self):
  1727. self.remLabels(self.canvas.deleteSelected())
  1728. self.actions.undo.setEnabled(True)
  1729. self.setDirty()
  1730. if self.noShapes():
  1731. for action in self.actions.onShapesPresent:
  1732. action.setEnabled(False)
  1733. self.BoxListDock.setWindowTitle(self.BoxListDockName + f" ({self.BoxList.count()})")
  1734. self.labelListDock.setWindowTitle(self.labelListDockName + f" ({self.labelList.count()})")
  1735. def chshapeLineColor(self):
  1736. color = self.colorDialog.getColor(self.lineColor, u'Choose line color',
  1737. default=DEFAULT_LINE_COLOR)
  1738. if color:
  1739. for shape in self.canvas.selectedShapes: shape.line_color = color
  1740. self.canvas.update()
  1741. self.setDirty()
  1742. def chshapeFillColor(self):
  1743. color = self.colorDialog.getColor(self.fillColor, u'Choose fill color',
  1744. default=DEFAULT_FILL_COLOR)
  1745. if color:
  1746. for shape in self.canvas.selectedShapes: shape.fill_color = color
  1747. self.canvas.update()
  1748. self.setDirty()
  1749. def copyShape(self):
  1750. self.canvas.endMove(copy=True)
  1751. self.addLabel(self.canvas.selectedShape)
  1752. self.setDirty()
  1753. def moveShape(self):
  1754. self.canvas.endMove(copy=False)
  1755. self.setDirty()
  1756. def loadPredefinedClasses(self, predefClassesFile):
  1757. if os.path.exists(predefClassesFile) is True:
  1758. with codecs.open(predefClassesFile, 'r', 'utf8') as f:
  1759. for line in f:
  1760. line = line.strip()
  1761. if self.labelHist is None:
  1762. self.labelHist = [line]
  1763. else:
  1764. self.labelHist.append(line)
  1765. def togglePaintLabelsOption(self):
  1766. self.displayIndexOption.setChecked(False)
  1767. for shape in self.canvas.shapes:
  1768. shape.paintLabel = self.displayLabelOption.isChecked()
  1769. shape.paintIdx = self.displayIndexOption.isChecked()
  1770. self.canvas.repaint()
  1771. def togglePaintIndexOption(self):
  1772. self.displayLabelOption.setChecked(False)
  1773. for shape in self.canvas.shapes:
  1774. shape.paintLabel = self.displayLabelOption.isChecked()
  1775. shape.paintIdx = self.displayIndexOption.isChecked()
  1776. self.canvas.repaint()
  1777. def toogleDrawSquare(self):
  1778. self.canvas.setDrawingShapeToSquare(self.drawSquaresOption.isChecked())
  1779. def additems(self, dirpath):
  1780. for file in self.mImgList:
  1781. pix = QPixmap(file)
  1782. _, filename = os.path.split(file)
  1783. filename, _ = os.path.splitext(filename)
  1784. item = QListWidgetItem(QIcon(pix.scaled(100, 100, Qt.IgnoreAspectRatio, Qt.FastTransformation)),
  1785. filename[:10])
  1786. item.setToolTip(file)
  1787. self.iconlist.addItem(item)
  1788. def additems5(self, dirpath):
  1789. for file in self.mImgList5:
  1790. pix = QPixmap(file)
  1791. _, filename = os.path.split(file)
  1792. filename, _ = os.path.splitext(filename)
  1793. pfilename = filename[:10]
  1794. if len(pfilename) < 10:
  1795. lentoken = 12 - len(pfilename)
  1796. prelen = lentoken // 2
  1797. bfilename = prelen * " " + pfilename + (lentoken - prelen) * " "
  1798. # item = QListWidgetItem(QIcon(pix.scaled(100, 100, Qt.KeepAspectRatio, Qt.SmoothTransformation)),filename[:10])
  1799. item = QListWidgetItem(QIcon(pix.scaled(100, 100, Qt.IgnoreAspectRatio, Qt.FastTransformation)), pfilename)
  1800. # item.setForeground(QBrush(Qt.white))
  1801. item.setToolTip(file)
  1802. self.iconlist.addItem(item)
  1803. owidth = 0
  1804. for index in range(len(self.mImgList5)):
  1805. item = self.iconlist.item(index)
  1806. itemwidget = self.iconlist.visualItemRect(item)
  1807. owidth += itemwidget.width()
  1808. self.iconlist.setMinimumWidth(owidth + 50)
  1809. def gen_quad_from_poly(self, poly):
  1810. """
  1811. Generate min area quad from poly.
  1812. """
  1813. point_num = poly.shape[0]
  1814. min_area_quad = np.zeros((4, 2), dtype=np.float32)
  1815. rect = cv2.minAreaRect(poly.astype(
  1816. np.int32)) # (center (x,y), (width, height), angle of rotation)
  1817. box = np.array(cv2.boxPoints(rect))
  1818. first_point_idx = 0
  1819. min_dist = 1e4
  1820. for i in range(4):
  1821. dist = np.linalg.norm(box[(i + 0) % 4] - poly[0]) + \
  1822. np.linalg.norm(box[(i + 1) % 4] - poly[point_num // 2 - 1]) + \
  1823. np.linalg.norm(box[(i + 2) % 4] - poly[point_num // 2]) + \
  1824. np.linalg.norm(box[(i + 3) % 4] - poly[-1])
  1825. if dist < min_dist:
  1826. min_dist = dist
  1827. first_point_idx = i
  1828. for i in range(4):
  1829. min_area_quad[i] = box[(first_point_idx + i) % 4]
  1830. bbox_new = min_area_quad.tolist()
  1831. bbox = []
  1832. for box in bbox_new:
  1833. box = list(map(int, box))
  1834. bbox.append(box)
  1835. return bbox
  1836. def getImglabelidx(self, filePath):
  1837. if platform.system() == 'Windows':
  1838. spliter = '\\'
  1839. else:
  1840. spliter = '/'
  1841. filepathsplit = filePath.split(spliter)[-2:]
  1842. return filepathsplit[0] + '/' + filepathsplit[1]
  1843. def autoRecognition(self):
  1844. assert self.mImgList is not None
  1845. print('Using model from ', self.model)
  1846. uncheckedList = [i for i in self.mImgList if i not in self.fileStatedict.keys()]
  1847. self.autoDialog = AutoDialog(parent=self, ocr=self.ocr, mImgList=uncheckedList, lenbar=len(uncheckedList))
  1848. self.autoDialog.popUp()
  1849. self.currIndex = len(self.mImgList) - 1
  1850. self.loadFile(self.filePath) # ADD
  1851. self.haveAutoReced = True
  1852. self.AutoRecognition.setEnabled(False)
  1853. self.actions.AutoRec.setEnabled(False)
  1854. self.setDirty()
  1855. self.saveCacheLabel()
  1856. self.init_key_list(self.Cachelabel)
  1857. def reRecognition(self):
  1858. img = cv2.imdecode(np.fromfile(self.filePath,dtype=np.uint8),1)
  1859. # org_box = [dic['points'] for dic in self.PPlabel[self.getImglabelidx(self.filePath)]]
  1860. if self.canvas.shapes:
  1861. self.result_dic = []
  1862. self.result_dic_locked = [] # result_dic_locked stores the ocr result of self.canvas.lockedShapes
  1863. rec_flag = 0
  1864. for shape in self.canvas.shapes:
  1865. box = [[int(p.x()), int(p.y())] for p in shape.points]
  1866. kie_cls = shape.key_cls
  1867. if len(box) > 4:
  1868. box = self.gen_quad_from_poly(np.array(box))
  1869. assert len(box) == 4
  1870. img_crop = get_rotate_crop_image(img, np.array(box, np.float32))
  1871. if img_crop is None:
  1872. msg = 'Can not recognise the detection box in ' + self.filePath + '. Please change manually'
  1873. QMessageBox.information(self, "Information", msg)
  1874. return
  1875. result = self.ocr.ocr(img_crop, cls=True, det=False)[0]
  1876. if result[0][0] != '':
  1877. if shape.line_color == DEFAULT_LOCK_COLOR:
  1878. shape.label = result[0][0]
  1879. result.insert(0, box)
  1880. if self.kie_mode:
  1881. result.append(kie_cls)
  1882. self.result_dic_locked.append(result)
  1883. else:
  1884. result.insert(0, box)
  1885. if self.kie_mode:
  1886. result.append(kie_cls)
  1887. self.result_dic.append(result)
  1888. else:
  1889. print('Can not recognise the box')
  1890. if shape.line_color == DEFAULT_LOCK_COLOR:
  1891. shape.label = result[0][0]
  1892. if self.kie_mode:
  1893. self.result_dic_locked.append([box, (self.noLabelText, 0), kie_cls])
  1894. else:
  1895. self.result_dic_locked.append([box, (self.noLabelText, 0)])
  1896. else:
  1897. if self.kie_mode:
  1898. self.result_dic.append([box, (self.noLabelText, 0), kie_cls])
  1899. else:
  1900. self.result_dic.append([box, (self.noLabelText, 0)])
  1901. try:
  1902. if self.noLabelText == shape.label or result[1][0] == shape.label:
  1903. print('label no change')
  1904. else:
  1905. rec_flag += 1
  1906. except IndexError as e:
  1907. print('Can not recognise the box')
  1908. if (len(self.result_dic) > 0 and rec_flag > 0) or self.canvas.lockedShapes:
  1909. self.canvas.isInTheSameImage = True
  1910. self.saveFile(mode='Auto')
  1911. self.loadFile(self.filePath)
  1912. self.canvas.isInTheSameImage = False
  1913. self.setDirty()
  1914. elif len(self.result_dic) == len(self.canvas.shapes) and rec_flag == 0:
  1915. if self.lang == 'ch':
  1916. QMessageBox.information(self, "Information", "识别结果保持一致!")
  1917. else:
  1918. QMessageBox.information(self, "Information", "The recognition result remains unchanged!")
  1919. else:
  1920. print('Can not recgonise in ', self.filePath)
  1921. else:
  1922. QMessageBox.information(self, "Information", "Draw a box!")
  1923. def singleRerecognition(self):
  1924. img = cv2.imdecode(np.fromfile(self.filePath,dtype=np.uint8),1)
  1925. for shape in self.canvas.selectedShapes:
  1926. box = [[int(p.x()), int(p.y())] for p in shape.points]
  1927. if len(box) > 4:
  1928. box = self.gen_quad_from_poly(np.array(box))
  1929. assert len(box) == 4
  1930. img_crop = get_rotate_crop_image(img, np.array(box, np.float32))
  1931. if img_crop is None:
  1932. msg = 'Can not recognise the detection box in ' + self.filePath + '. Please change manually'
  1933. QMessageBox.information(self, "Information", msg)
  1934. return
  1935. result = self.ocr.ocr(img_crop, cls=True, det=False)[0]
  1936. if result[0][0] != '':
  1937. result.insert(0, box)
  1938. print('result in reRec is ', result)
  1939. if result[1][0] == shape.label:
  1940. print('label no change')
  1941. else:
  1942. shape.label = result[1][0]
  1943. else:
  1944. print('Can not recognise the box')
  1945. if self.noLabelText == shape.label:
  1946. print('label no change')
  1947. else:
  1948. shape.label = self.noLabelText
  1949. self.singleLabel(shape)
  1950. self.setDirty()
  1951. def TableRecognition(self):
  1952. '''
  1953. Table Recegnition
  1954. '''
  1955. from paddleocr import to_excel
  1956. import time
  1957. start = time.time()
  1958. img = cv2.imread(self.filePath)
  1959. res = self.table_ocr(img, return_ocr_result_in_table=True)
  1960. TableRec_excel_dir = self.lastOpenDir + '/tableRec_excel_output/'
  1961. os.makedirs(TableRec_excel_dir, exist_ok=True)
  1962. filename, _ = os.path.splitext(os.path.basename(self.filePath))
  1963. excel_path = TableRec_excel_dir + '{}.xlsx'.format(filename)
  1964. if res is None:
  1965. msg = 'Can not recognise the table in ' + self.filePath + '. Please change manually'
  1966. QMessageBox.information(self, "Information", msg)
  1967. to_excel('', excel_path) # create an empty excel
  1968. return
  1969. # save res
  1970. # ONLY SUPPORT ONE TABLE in one image
  1971. hasTable = False
  1972. for region in res:
  1973. if region['type'] == 'table':
  1974. if region['res']['boxes'] is None:
  1975. msg = 'Can not recognise the detection box in ' + self.filePath + '. Please change manually'
  1976. QMessageBox.information(self, "Information", msg)
  1977. to_excel('', excel_path) # create an empty excel
  1978. return
  1979. hasTable = True
  1980. # save table ocr result on PPOCRLabel
  1981. # clear all old annotaions before saving result
  1982. self.itemsToShapes.clear()
  1983. self.shapesToItems.clear()
  1984. self.itemsToShapesbox.clear() # ADD
  1985. self.shapesToItemsbox.clear()
  1986. self.labelList.clear()
  1987. self.indexList.clear()
  1988. self.BoxList.clear()
  1989. self.result_dic = []
  1990. self.result_dic_locked = []
  1991. shapes = []
  1992. result_len = len(region['res']['boxes'])
  1993. order_index = 0
  1994. for i in range(result_len):
  1995. bbox = np.array(region['res']['boxes'][i])
  1996. rec_text = region['res']['rec_res'][i][0]
  1997. rext_bbox = [[bbox[0], bbox[1]], [bbox[2], bbox[1]], [bbox[2], bbox[3]], [bbox[0], bbox[3]]]
  1998. # save bbox to shape
  1999. shape = Shape(label=rec_text, line_color=DEFAULT_LINE_COLOR, key_cls=None)
  2000. for point in rext_bbox:
  2001. x, y = point
  2002. # Ensure the labels are within the bounds of the image.
  2003. # If not, fix them.
  2004. x, y, snapped = self.canvas.snapPointToCanvas(x, y)
  2005. shape.addPoint(QPointF(x, y))
  2006. shape.difficult = False
  2007. shape.idx = order_index
  2008. order_index += 1
  2009. # shape.locked = False
  2010. shape.close()
  2011. self.addLabel(shape)
  2012. shapes.append(shape)
  2013. self.setDirty()
  2014. self.canvas.loadShapes(shapes)
  2015. # save HTML result to excel
  2016. try:
  2017. to_excel(region['res']['html'], excel_path)
  2018. except:
  2019. print('Can not save excel file, maybe Permission denied (.xlsx is being occupied)')
  2020. break
  2021. if not hasTable:
  2022. msg = 'Can not recognise the table in ' + self.filePath + '. Please change manually'
  2023. QMessageBox.information(self, "Information", msg)
  2024. to_excel('', excel_path) # create an empty excel
  2025. return
  2026. # automatically open excel annotation file
  2027. if platform.system() == 'Windows':
  2028. try:
  2029. import win32com.client
  2030. except:
  2031. print("CANNOT OPEN .xlsx. It could be one of the following reasons: " \
  2032. "Only support Windows | No python win32com")
  2033. try:
  2034. xl = win32com.client.Dispatch("Excel.Application")
  2035. xl.Visible = True
  2036. xl.Workbooks.Open(excel_path)
  2037. # excelEx = "You need to show the excel executable at this point"
  2038. # subprocess.Popen([excelEx, excel_path])
  2039. # os.startfile(excel_path)
  2040. except:
  2041. print("CANNOT OPEN .xlsx. It could be the following reasons: " \
  2042. ".xlsx is not existed")
  2043. else:
  2044. os.system('open ' + os.path.normpath(excel_path))
  2045. print('time cost: ', time.time() - start)
  2046. def cellreRecognition(self):
  2047. '''
  2048. re-recognise text in a cell
  2049. '''
  2050. img = cv2.imread(self.filePath)
  2051. for shape in self.canvas.selectedShapes:
  2052. box = [[int(p.x()), int(p.y())] for p in shape.points]
  2053. if len(box) > 4:
  2054. box = self.gen_quad_from_poly(np.array(box))
  2055. assert len(box) == 4
  2056. # pad around bbox for better text recognition accuracy
  2057. _box = boxPad(box, img.shape, 6)
  2058. img_crop = get_rotate_crop_image(img, np.array(_box, np.float32))
  2059. if img_crop is None:
  2060. msg = 'Can not recognise the detection box in ' + self.filePath + '. Please change manually'
  2061. QMessageBox.information(self, "Information", msg)
  2062. return
  2063. # merge the text result in the cell
  2064. texts = ''
  2065. probs = 0. # the probability of the cell is avgerage prob of every text box in the cell
  2066. bboxes = self.ocr.ocr(img_crop, det=True, rec=False, cls=False)[0]
  2067. if len(bboxes) > 0:
  2068. bboxes.reverse() # top row text at first
  2069. for _bbox in bboxes:
  2070. patch = get_rotate_crop_image(img_crop, np.array(_bbox, np.float32))
  2071. rec_res = self.ocr.ocr(patch, det=False, rec=True, cls=False)[0]
  2072. text = rec_res[0][0]
  2073. if text != '':
  2074. texts += text + ('' if text[0].isalpha() else ' ') # add space between english word
  2075. probs += rec_res[0][1]
  2076. probs = probs / len(bboxes)
  2077. result = [(texts.strip(), probs)]
  2078. if result[0][0] != '':
  2079. result.insert(0, box)
  2080. print('result in reRec is ', result)
  2081. if result[1][0] == shape.label:
  2082. print('label no change')
  2083. else:
  2084. shape.label = result[1][0]
  2085. else:
  2086. print('Can not recognise the box')
  2087. if self.noLabelText == shape.label:
  2088. print('label no change')
  2089. else:
  2090. shape.label = self.noLabelText
  2091. self.singleLabel(shape)
  2092. self.setDirty()
  2093. def exportJSON(self):
  2094. '''
  2095. export PPLabel and CSV to JSON (PubTabNet)
  2096. '''
  2097. import pandas as pd
  2098. # automatically save annotations
  2099. self.saveFilestate()
  2100. self.savePPlabel(mode='auto')
  2101. # load box annotations
  2102. labeldict = {}
  2103. if not os.path.exists(self.PPlabelpath):
  2104. msg = 'ERROR, Can not find Label.txt'
  2105. QMessageBox.information(self, "Information", msg)
  2106. return
  2107. else:
  2108. with open(self.PPlabelpath, 'r', encoding='utf-8') as f:
  2109. data = f.readlines()
  2110. for each in data:
  2111. file, label = each.split('\t')
  2112. if label:
  2113. label = label.replace('false', 'False')
  2114. label = label.replace('true', 'True')
  2115. labeldict[file] = eval(label)
  2116. else:
  2117. labeldict[file] = []
  2118. # read table recognition output
  2119. TableRec_excel_dir = os.path.join(
  2120. self.lastOpenDir, 'tableRec_excel_output')
  2121. # save txt
  2122. fid = open(
  2123. "{}/gt.txt".format(self.lastOpenDir), "w", encoding='utf-8')
  2124. for image_path in labeldict.keys():
  2125. # load csv annotations
  2126. filename, _ = os.path.splitext(os.path.basename(image_path))
  2127. csv_path = os.path.join(
  2128. TableRec_excel_dir, filename + '.xlsx')
  2129. if not os.path.exists(csv_path):
  2130. continue
  2131. excel = xlrd.open_workbook(csv_path)
  2132. sheet0 = excel.sheet_by_index(0) # only sheet 0
  2133. merged_cells = sheet0.merged_cells # (0,1,1,3) start row, end row, start col, end col
  2134. html_list = [['td'] * sheet0.ncols for i in range(sheet0.nrows)]
  2135. for merged in merged_cells:
  2136. html_list = expand_list(merged, html_list)
  2137. token_list = convert_token(html_list)
  2138. # load box annotations
  2139. cells = []
  2140. for anno in labeldict[image_path]:
  2141. tokens = list(anno['transcription'])
  2142. cells.append({
  2143. 'tokens': tokens,
  2144. 'bbox': anno['points']
  2145. })
  2146. # 构造标注信息
  2147. html = {
  2148. 'structure': {
  2149. 'tokens': token_list
  2150. },
  2151. 'cells': cells
  2152. }
  2153. d = {
  2154. 'filename': os.path.basename(image_path),
  2155. 'html': html
  2156. }
  2157. # 重构HTML
  2158. d['gt'] = rebuild_html_from_ppstructure_label(d)
  2159. fid.write('{}\n'.format(
  2160. json.dumps(
  2161. d, ensure_ascii=False)))
  2162. # convert to PP-Structure label format
  2163. fid.close()
  2164. msg = 'JSON sucessfully saved in {}/gt.txt'.format(self.lastOpenDir)
  2165. QMessageBox.information(self, "Information", msg)
  2166. def autolcm(self):
  2167. vbox = QVBoxLayout()
  2168. hbox = QHBoxLayout()
  2169. self.panel = QLabel()
  2170. self.panel.setText(self.stringBundle.getString('choseModelLg'))
  2171. self.panel.setAlignment(Qt.AlignLeft)
  2172. self.comboBox = QComboBox()
  2173. self.comboBox.setObjectName("comboBox")
  2174. self.comboBox.addItems(['Chinese & English', 'English', 'French', 'German', 'Korean', 'Japanese'])
  2175. vbox.addWidget(self.panel)
  2176. vbox.addWidget(self.comboBox)
  2177. self.dialog = QDialog()
  2178. self.dialog.resize(300, 100)
  2179. self.okBtn = QPushButton(self.stringBundle.getString('ok'))
  2180. self.cancelBtn = QPushButton(self.stringBundle.getString('cancel'))
  2181. self.okBtn.clicked.connect(self.modelChoose)
  2182. self.cancelBtn.clicked.connect(self.cancel)
  2183. self.dialog.setWindowTitle(self.stringBundle.getString('choseModelLg'))
  2184. hbox.addWidget(self.okBtn)
  2185. hbox.addWidget(self.cancelBtn)
  2186. vbox.addWidget(self.panel)
  2187. vbox.addLayout(hbox)
  2188. self.dialog.setLayout(vbox)
  2189. self.dialog.setWindowModality(Qt.ApplicationModal)
  2190. self.dialog.exec_()
  2191. if self.filePath:
  2192. self.AutoRecognition.setEnabled(True)
  2193. self.actions.AutoRec.setEnabled(True)
  2194. def modelChoose(self):
  2195. print(self.comboBox.currentText())
  2196. lg_idx = {'Chinese & English': 'ch', 'English': 'en', 'French': 'french', 'German': 'german',
  2197. 'Korean': 'korean', 'Japanese': 'japan'}
  2198. del self.ocr
  2199. self.ocr = PaddleOCR(use_pdserving=False, use_angle_cls=True, det=True, cls=True, use_gpu=False,
  2200. lang=lg_idx[self.comboBox.currentText()])
  2201. del self.table_ocr
  2202. self.table_ocr = PPStructure(use_pdserving=False,
  2203. use_gpu=False,
  2204. lang=lg_idx[self.comboBox.currentText()],
  2205. layout=False,
  2206. show_log=False)
  2207. self.dialog.close()
  2208. def cancel(self):
  2209. self.dialog.close()
  2210. def loadFilestate(self, saveDir):
  2211. self.fileStatepath = saveDir + '/fileState.txt'
  2212. self.fileStatedict = {}
  2213. if not os.path.exists(self.fileStatepath):
  2214. f = open(self.fileStatepath, 'w', encoding='utf-8')
  2215. else:
  2216. with open(self.fileStatepath, 'r', encoding='utf-8') as f:
  2217. states = f.readlines()
  2218. for each in states:
  2219. file, state = each.split('\t')
  2220. self.fileStatedict[file] = 1
  2221. self.actions.saveLabel.setEnabled(True)
  2222. self.actions.saveRec.setEnabled(True)
  2223. self.actions.exportJSON.setEnabled(True)
  2224. def saveFilestate(self):
  2225. with open(self.fileStatepath, 'w', encoding='utf-8') as f:
  2226. for key in self.fileStatedict:
  2227. f.write(key + '\t')
  2228. f.write(str(self.fileStatedict[key]) + '\n')
  2229. def loadLabelFile(self, labelpath):
  2230. labeldict = {}
  2231. if not os.path.exists(labelpath):
  2232. f = open(labelpath, 'w', encoding='utf-8')
  2233. else:
  2234. with open(labelpath, 'r', encoding='utf-8') as f:
  2235. data = f.readlines()
  2236. for each in data:
  2237. file, label = each.split('\t')
  2238. if label:
  2239. label = label.replace('false', 'False')
  2240. label = label.replace('true', 'True')
  2241. labeldict[file] = eval(label)
  2242. else:
  2243. labeldict[file] = []
  2244. return labeldict
  2245. def savePPlabel(self, mode='Manual'):
  2246. savedfile = [self.getImglabelidx(i) for i in self.fileStatedict.keys()]
  2247. with open(self.PPlabelpath, 'w', encoding='utf-8') as f:
  2248. for key in self.PPlabel:
  2249. if key in savedfile and self.PPlabel[key] != []:
  2250. f.write(key + '\t')
  2251. f.write(json.dumps(self.PPlabel[key], ensure_ascii=False) + '\n')
  2252. if mode == 'Manual':
  2253. if self.lang == 'ch':
  2254. msg = '已将检查过的图片标签保存在 ' + self.PPlabelpath + " 文件中"
  2255. else:
  2256. msg = 'Images that have been checked are saved in ' + self.PPlabelpath
  2257. QMessageBox.information(self, "Information", msg)
  2258. def saveCacheLabel(self):
  2259. with open(self.Cachelabelpath, 'w', encoding='utf-8') as f:
  2260. for key in self.Cachelabel:
  2261. f.write(key + '\t')
  2262. f.write(json.dumps(self.Cachelabel[key], ensure_ascii=False) + '\n')
  2263. def saveLabelFile(self):
  2264. self.saveFilestate()
  2265. self.savePPlabel()
  2266. def saveRecResult(self):
  2267. if {} in [self.PPlabelpath, self.PPlabel, self.fileStatedict]:
  2268. QMessageBox.information(self, "Information", "Check the image first")
  2269. return
  2270. rec_gt_dir = os.path.dirname(self.PPlabelpath) + '/rec_gt.txt'
  2271. crop_img_dir = os.path.dirname(self.PPlabelpath) + '/crop_img/'
  2272. ques_img = []
  2273. if not os.path.exists(crop_img_dir):
  2274. os.mkdir(crop_img_dir)
  2275. with open(rec_gt_dir, 'w', encoding='utf-8') as f:
  2276. for key in self.fileStatedict:
  2277. idx = self.getImglabelidx(key)
  2278. try:
  2279. img = cv2.imread(key)
  2280. for i, label in enumerate(self.PPlabel[idx]):
  2281. if label['difficult']:
  2282. continue
  2283. img_crop = get_rotate_crop_image(img, np.array(label['points'], np.float32))
  2284. img_name = os.path.splitext(os.path.basename(idx))[0] + '_crop_' + str(i) + '.jpg'
  2285. cv2.imwrite(crop_img_dir + img_name, img_crop)
  2286. f.write('crop_img/' + img_name + '\t')
  2287. f.write(label['transcription'] + '\n')
  2288. except Exception as e:
  2289. ques_img.append(key)
  2290. print("Can not read image ", e)
  2291. if ques_img:
  2292. QMessageBox.information(self,
  2293. "Information",
  2294. "The following images can not be saved, please check the image path and labels.\n"
  2295. + "".join(str(i) + '\n' for i in ques_img))
  2296. QMessageBox.information(self, "Information", "Cropped images have been saved in " + str(crop_img_dir))
  2297. def speedChoose(self):
  2298. if self.labelDialogOption.isChecked():
  2299. self.canvas.newShape.disconnect()
  2300. self.canvas.newShape.connect(partial(self.newShape, True))
  2301. else:
  2302. self.canvas.newShape.disconnect()
  2303. self.canvas.newShape.connect(partial(self.newShape, False))
  2304. def autoSaveFunc(self):
  2305. if self.autoSaveOption.isChecked():
  2306. self.autoSaveNum = 1 # Real auto_Save
  2307. try:
  2308. self.saveLabelFile()
  2309. except:
  2310. pass
  2311. print('The program will automatically save once after confirming an image')
  2312. else:
  2313. self.autoSaveNum = 5 # Used for backup
  2314. print('The program will automatically save once after confirming 5 images (default)')
  2315. def change_box_key(self):
  2316. if not self.kie_mode:
  2317. return
  2318. key_text, _ = self.keyDialog.popUp(self.key_previous_text)
  2319. if key_text is None:
  2320. return
  2321. self.key_previous_text = key_text
  2322. for shape in self.canvas.selectedShapes:
  2323. shape.key_cls = key_text
  2324. if not self.keyList.findItemsByLabel(key_text):
  2325. item = self.keyList.createItemFromLabel(key_text)
  2326. self.keyList.addItem(item)
  2327. rgb = self._get_rgb_by_label(key_text, self.kie_mode)
  2328. self.keyList.setItemLabel(item, key_text, rgb)
  2329. self._update_shape_color(shape)
  2330. self.keyDialog.addLabelHistory(key_text)
  2331. # save changed shape
  2332. self.setDirty()
  2333. def undoShapeEdit(self):
  2334. self.canvas.restoreShape()
  2335. self.labelList.clear()
  2336. self.indexList.clear()
  2337. self.BoxList.clear()
  2338. self.loadShapes(self.canvas.shapes)
  2339. self.actions.undo.setEnabled(self.canvas.isShapeRestorable)
  2340. def loadShapes(self, shapes, replace=True):
  2341. self._noSelectionSlot = True
  2342. for shape in shapes:
  2343. self.addLabel(shape)
  2344. self.labelList.clearSelection()
  2345. self.indexList.clearSelection()
  2346. self._noSelectionSlot = False
  2347. self.canvas.loadShapes(shapes, replace=replace)
  2348. print("loadShapes") # 1
  2349. def lockSelectedShape(self):
  2350. """lock the selected shapes.
  2351. Add self.selectedShapes to lock self.canvas.lockedShapes,
  2352. which holds the ratio of the four coordinates of the locked shapes
  2353. to the width and height of the image
  2354. """
  2355. width, height = self.image.width(), self.image.height()
  2356. def format_shape(s):
  2357. return dict(label=s.label, # str
  2358. line_color=s.line_color.getRgb(),
  2359. fill_color=s.fill_color.getRgb(),
  2360. ratio=[[int(p.x()) / width, int(p.y()) / height] for p in s.points], # QPonitF
  2361. difficult=s.difficult, # bool
  2362. key_cls=s.key_cls, # bool
  2363. )
  2364. # lock
  2365. if len(self.canvas.lockedShapes) == 0:
  2366. for s in self.canvas.selectedShapes:
  2367. s.line_color = DEFAULT_LOCK_COLOR
  2368. s.locked = True
  2369. shapes = [format_shape(shape) for shape in self.canvas.selectedShapes]
  2370. trans_dic = []
  2371. for box in shapes:
  2372. trans_dict = {"transcription": box['label'], "ratio": box['ratio'], "difficult": box['difficult']}
  2373. if self.kie_mode:
  2374. trans_dict.update({"key_cls": box["key_cls"]})
  2375. trans_dic.append(trans_dict)
  2376. self.canvas.lockedShapes = trans_dic
  2377. self.actions.save.setEnabled(True)
  2378. # unlock
  2379. else:
  2380. for s in self.canvas.shapes:
  2381. s.line_color = DEFAULT_LINE_COLOR
  2382. self.canvas.lockedShapes = []
  2383. self.result_dic_locked = []
  2384. self.setDirty()
  2385. self.actions.save.setEnabled(True)
  2386. def inverted(color):
  2387. return QColor(*[255 - v for v in color.getRgb()])
  2388. def read(filename, default=None):
  2389. try:
  2390. with open(filename, 'rb') as f:
  2391. return f.read()
  2392. except:
  2393. return default
  2394. def str2bool(v):
  2395. return v.lower() in ("true", "t", "1")
  2396. def get_main_app(argv=[]):
  2397. """
  2398. Standard boilerplate Qt application code.
  2399. Do everything but app.exec_() -- so that we can test the application in one thread
  2400. """
  2401. app = QApplication(argv)
  2402. app.setApplicationName(__appname__)
  2403. app.setWindowIcon(newIcon("app"))
  2404. # Tzutalin 201705+: Accept extra arguments to change predefined class file
  2405. arg_parser = argparse.ArgumentParser()
  2406. arg_parser.add_argument("--lang", type=str, default='en', nargs="?")
  2407. arg_parser.add_argument("--gpu", type=str2bool, default=True, nargs="?")
  2408. arg_parser.add_argument("--kie", type=str2bool, default=False, nargs="?")
  2409. arg_parser.add_argument("--predefined_classes_file",
  2410. default=os.path.join(os.path.dirname(__file__), "data", "predefined_classes.txt"),
  2411. nargs="?")
  2412. args = arg_parser.parse_args(argv[1:])
  2413. win = MainWindow(lang=args.lang,
  2414. gpu=args.gpu,
  2415. kie_mode=args.kie,
  2416. default_predefined_class_file=args.predefined_classes_file)
  2417. win.show()
  2418. return app, win
  2419. def main():
  2420. """construct main app and run it"""
  2421. app, _win = get_main_app(sys.argv)
  2422. return app.exec_()
  2423. if __name__ == '__main__':
  2424. resource_file = './libs/resources.py'
  2425. if not os.path.exists(resource_file):
  2426. output = os.system('pyrcc5 -o libs/resources.py resources.qrc')
  2427. assert output == 0, "operate the cmd have some problems ,please check whether there is a in the lib " \
  2428. "directory resources.py "
  2429. sys.exit(main())