1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747174817491750175117521753175417551756175717581759176017611762176317641765176617671768176917701771177217731774177517761777177817791780178117821783178417851786178717881789179017911792179317941795179617971798179918001801180218031804180518061807180818091810181118121813181418151816181718181819182018211822182318241825182618271828182918301831183218331834183518361837183818391840184118421843184418451846184718481849185018511852185318541855185618571858185918601861186218631864186518661867186818691870187118721873187418751876187718781879188018811882188318841885188618871888188918901891189218931894189518961897189818991900190119021903190419051906190719081909191019111912191319141915191619171918191919201921192219231924192519261927192819291930193119321933193419351936193719381939194019411942194319441945194619471948194919501951195219531954195519561957195819591960196119621963196419651966196719681969197019711972197319741975197619771978197919801981198219831984198519861987198819891990199119921993199419951996199719981999200020012002200320042005 |
- # Copyright (c) <2015-Present> Tzutalin
- # Copyright (C) 2013 MIT, Computer Science and Artificial Intelligence Laboratory. Bryan Russell, Antonio Torralba,
- # William T. Freeman. Permission is hereby granted, free of charge, to any person obtaining a copy of this software and
- # associated documentation files (the "Software"), to deal in the Software without restriction, including without
- # limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the
- # Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
- # The above copyright notice and this permission notice shall be included in all copies or substantial portions of
- # the Software. THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT
- # NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
- # SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF
- # CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
- # THE SOFTWARE.
- #!/usr/bin/env python
- # -*- coding: utf-8 -*-
- # pyrcc5 -o libs/resources.py resources.qrc
- import argparse
- import ast
- import codecs
- import os.path
- import platform
- import subprocess
- import sys
- from functools import partial
- from collections import defaultdict
- import json
- import cv2
- __dir__ = os.path.dirname(os.path.abspath(__file__))
- sys.path.append(__dir__)
- sys.path.append(os.path.abspath(os.path.join(__dir__, '../..')))
- sys.path.append("..")
- from paddleocr import PaddleOCR
- try:
- from PyQt5 import QtCore, QtGui, QtWidgets
- from PyQt5.QtGui import *
- from PyQt5.QtCore import *
- from PyQt5.QtWidgets import *
- except ImportError:
- # needed for py3+qt4
- # Ref:
- # http://pyqt.sourceforge.net/Docs/PyQt4/incompatible_apis.html
- # http://stackoverflow.com/questions/21217399/pyqt4-qtcore-qvariant-object-instead-of-a-string
- if sys.version_info.major >= 3:
- import sip
- sip.setapi('QVariant', 2)
- from PyQt4.QtGui import *
- from PyQt4.QtCore import *
- from combobox import ComboBox
- from libs.constants import *
- from libs.utils import *
- from libs.settings import Settings
- from libs.shape import Shape, DEFAULT_LINE_COLOR, DEFAULT_FILL_COLOR
- from libs.stringBundle import StringBundle
- from libs.canvas import Canvas
- from libs.zoomWidget import ZoomWidget
- from libs.autoDialog import AutoDialog
- from libs.labelDialog import LabelDialog
- from libs.colorDialog import ColorDialog
- from libs.toolBar import ToolBar
- from libs.ustr import ustr
- from libs.hashableQListWidgetItem import HashableQListWidgetItem
- from libs.editinlist import EditInList
- __appname__ = 'PPOCRLabel'
- class WindowMixin(object):
- def menu(self, title, actions=None):
- menu = self.menuBar().addMenu(title)
- if actions:
- addActions(menu, actions)
- return menu
- def toolbar(self, title, actions=None):
- toolbar = ToolBar(title)
- toolbar.setObjectName(u'%sToolBar' % title)
- # toolbar.setOrientation(Qt.Vertical)
- toolbar.setToolButtonStyle(Qt.ToolButtonTextUnderIcon)
- if actions:
- addActions(toolbar, actions)
- self.addToolBar(Qt.LeftToolBarArea, toolbar)
- return toolbar
- class MainWindow(QMainWindow, WindowMixin):
- FIT_WINDOW, FIT_WIDTH, MANUAL_ZOOM = list(range(3))
- def __init__(self, lang="ch", defaultFilename=None, defaultPrefdefClassFile=None, defaultSaveDir=None):
- super(MainWindow, self).__init__()
- self.setWindowTitle(__appname__)
- # Load setting in the main thread
- self.settings = Settings()
- self.settings.load()
- settings = self.settings
- self.lang = lang
- # Load string bundle for i18n
- if lang not in ['ch', 'en']:
- lang = 'en'
- self.stringBundle = StringBundle.getBundle(localeStr='zh-CN' if lang=='ch' else 'en') # 'en'
- getStr = lambda strId: self.stringBundle.getString(strId)
- self.defaultSaveDir = defaultSaveDir
- self.ocr = PaddleOCR(use_pdserving=False, use_angle_cls=True, det=True, cls=True, use_gpu=False, lang=lang)
- if os.path.exists('./data/paddle.png'):
- result = self.ocr.ocr('./data/paddle.png', cls=True, det=True)
- # For loading all image under a directory
- self.mImgList = []
- self.mImgList5 = []
- self.dirname = None
- self.labelHist = []
- self.lastOpenDir = None
- self.result_dic = []
- self.changeFileFolder = False
- self.haveAutoReced = False
- self.labelFile = None
- self.currIndex = 0
- # Whether we need to save or not.
- self.dirty = False
- self._noSelectionSlot = False
- self._beginner = True
- self.screencastViewer = self.getAvailableScreencastViewer()
- self.screencast = "https://github.com/PaddlePaddle/PaddleOCR"
- # Load predefined classes to the list
- self.loadPredefinedClasses(defaultPrefdefClassFile)
- # Main widgets and related state.
- self.labelDialog = LabelDialog(parent=self, listItem=self.labelHist)
- self.autoDialog = AutoDialog(parent=self)
- self.itemsToShapes = {}
- self.shapesToItems = {}
- self.itemsToShapesbox = {}
- self.shapesToItemsbox = {}
- self.prevLabelText = getStr('tempLabel')
- self.model = 'paddle'
- self.PPreader = None
- self.autoSaveNum = 5
- ################# file list ###############
- self.fileListWidget = QListWidget()
- self.fileListWidget.itemClicked.connect(self.fileitemDoubleClicked)
- self.fileListWidget.setIconSize(QSize(25, 25))
- filelistLayout = QVBoxLayout()
- filelistLayout.setContentsMargins(0, 0, 0, 0)
- filelistLayout.addWidget(self.fileListWidget)
-
- self.AutoRecognition = QToolButton()
- self.AutoRecognition.setToolButtonStyle(Qt.ToolButtonTextBesideIcon)
- self.AutoRecognition.setIcon(newIcon('Auto'))
- autoRecLayout = QHBoxLayout()
- autoRecLayout.setContentsMargins(0, 0, 0, 0)
- autoRecLayout.addWidget(self.AutoRecognition)
- autoRecContainer = QWidget()
- autoRecContainer.setLayout(autoRecLayout)
- filelistLayout.addWidget(autoRecContainer)
- fileListContainer = QWidget()
- fileListContainer.setLayout(filelistLayout)
- self.filedock = QDockWidget(getStr('fileList'), self)
- self.filedock.setObjectName(getStr('files'))
- self.filedock.setWidget(fileListContainer)
- self.addDockWidget(Qt.LeftDockWidgetArea, self.filedock)
-
- ######## Right area ##########
- listLayout = QVBoxLayout()
- listLayout.setContentsMargins(0, 0, 0, 0)
- self.editButton = QToolButton()
- self.reRecogButton = QToolButton()
- self.reRecogButton.setIcon(newIcon('reRec', 30))
- self.reRecogButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon)
- self.newButton = QToolButton()
- self.newButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon)
- self.SaveButton = QToolButton()
- self.SaveButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon)
- self.DelButton = QToolButton()
- self.DelButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon)
- lefttoptoolbox = QHBoxLayout()
- lefttoptoolbox.addWidget(self.newButton)
- lefttoptoolbox.addWidget(self.reRecogButton)
- lefttoptoolboxcontainer = QWidget()
- lefttoptoolboxcontainer.setLayout(lefttoptoolbox)
- listLayout.addWidget(lefttoptoolboxcontainer)
- ################## label list ####################
- # Create and add a widget for showing current label items
- self.labelList = EditInList()
- labelListContainer = QWidget()
- labelListContainer.setLayout(listLayout)
- self.labelList.itemActivated.connect(self.labelSelectionChanged)
- self.labelList.itemSelectionChanged.connect(self.labelSelectionChanged)
- self.labelList.clicked.connect(self.labelList.item_clicked)
- # Connect to itemChanged to detect checkbox changes.
- self.labelList.itemChanged.connect(self.labelItemChanged)
- self.labelListDock = QDockWidget(getStr('recognitionResult'),self)
- self.labelListDock.setWidget(self.labelList)
- self.labelListDock.setFeatures(QDockWidget.NoDockWidgetFeatures)
- listLayout.addWidget(self.labelListDock)
- ################## detection box ####################
- self.BoxList = QListWidget()
- self.BoxList.itemActivated.connect(self.boxSelectionChanged)
- self.BoxList.itemSelectionChanged.connect(self.boxSelectionChanged)
- self.BoxList.itemDoubleClicked.connect(self.editBox)
- # Connect to itemChanged to detect checkbox changes.
- self.BoxList.itemChanged.connect(self.boxItemChanged)
- self.BoxListDock = QDockWidget(getStr('detectionBoxposition'), self)
- self.BoxListDock.setWidget(self.BoxList)
- self.BoxListDock.setFeatures(QDockWidget.NoDockWidgetFeatures)
- listLayout.addWidget(self.BoxListDock)
- ############ lower right area ############
- leftbtmtoolbox = QHBoxLayout()
- leftbtmtoolbox.addWidget(self.SaveButton)
- leftbtmtoolbox.addWidget(self.DelButton)
- leftbtmtoolboxcontainer = QWidget()
- leftbtmtoolboxcontainer.setLayout(leftbtmtoolbox)
- listLayout.addWidget(leftbtmtoolboxcontainer)
- self.dock = QDockWidget(getStr('boxLabelText'), self)
- self.dock.setObjectName(getStr('labels'))
- self.dock.setWidget(labelListContainer)
- ########## zoom bar #########
- self.imgsplider = QSlider(Qt.Horizontal)
- self.imgsplider.valueChanged.connect(self.CanvasSizeChange)
- self.imgsplider.setMinimum(-150)
- self.imgsplider.setMaximum(150)
- self.imgsplider.setSingleStep(1)
- self.imgsplider.setTickPosition(QSlider.TicksBelow)
- self.imgsplider.setTickInterval(1)
- op = QGraphicsOpacityEffect()
- op.setOpacity(0.2)
- self.imgsplider.setGraphicsEffect(op)
- # self.imgsplider.setAttribute(Qt.WA_TranslucentBackground)
- self.imgsplider.setStyleSheet("background-color:transparent")
- self.imgsliderDock = QDockWidget(getStr('ImageResize'), self)
- self.imgsliderDock.setObjectName(getStr('IR'))
- self.imgsliderDock.setWidget(self.imgsplider)
- self.imgsliderDock.setFeatures(QDockWidget.DockWidgetFloatable)
- self.imgsliderDock.setAttribute(Qt.WA_TranslucentBackground)
- self.addDockWidget(Qt.RightDockWidgetArea, self.imgsliderDock)
- self.zoomWidget = ZoomWidget()
- self.colorDialog = ColorDialog(parent=self)
- self.zoomWidgetValue = self.zoomWidget.value()
-
- ########## thumbnail #########
- hlayout = QHBoxLayout()
- m = (0, 0, 0, 0)
- hlayout.setSpacing(0)
- hlayout.setContentsMargins(*m)
- self.preButton = QToolButton()
- self.preButton.setIcon(newIcon("prev",40))
- self.preButton.setIconSize(QSize(40, 100))
- self.preButton.clicked.connect(self.openPrevImg)
- self.preButton.setStyleSheet('border: none;')
- self.preButton.setShortcut('a')
- self.iconlist = QListWidget()
- self.iconlist.setViewMode(QListView.IconMode)
- self.iconlist.setFlow(QListView.TopToBottom)
- self.iconlist.setSpacing(10)
- self.iconlist.setIconSize(QSize(50, 50))
- self.iconlist.setMovement(False)
- self.iconlist.setResizeMode(QListView.Adjust)
- self.iconlist.itemClicked.connect(self.iconitemDoubleClicked)
- self.iconlist.setStyleSheet("background-color:transparent; border: none;")
- self.iconlist.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff)
- self.nextButton = QToolButton()
- self.nextButton.setIcon(newIcon("next", 40))
- self.nextButton.setIconSize(QSize(40, 100))
- self.nextButton.setStyleSheet('border: none;')
- self.nextButton.clicked.connect(self.openNextImg)
- self.nextButton.setShortcut('d')
-
- hlayout.addWidget(self.preButton)
- hlayout.addWidget(self.iconlist)
- hlayout.addWidget(self.nextButton)
- iconListContainer = QWidget()
- iconListContainer.setLayout(hlayout)
- iconListContainer.setFixedHeight(100)
-
- ########### Canvas ###########
- self.canvas = Canvas(parent=self)
- self.canvas.zoomRequest.connect(self.zoomRequest)
- self.canvas.setDrawingShapeToSquare(settings.get(SETTING_DRAW_SQUARE, False))
- scroll = QScrollArea()
- scroll.setWidget(self.canvas)
- scroll.setWidgetResizable(True)
- self.scrollBars = {
- Qt.Vertical: scroll.verticalScrollBar(),
- Qt.Horizontal: scroll.horizontalScrollBar()
- }
- self.scrollArea = scroll
- self.canvas.scrollRequest.connect(self.scrollRequest)
- self.canvas.newShape.connect(partial(self.newShape, False))
- self.canvas.shapeMoved.connect(self.updateBoxlist) # self.setDirty
- self.canvas.selectionChanged.connect(self.shapeSelectionChanged)
- self.canvas.drawingPolygon.connect(self.toggleDrawingSensitive)
- centerLayout = QVBoxLayout()
- centerLayout.setContentsMargins(0, 0, 0, 0)
- centerLayout.addWidget(scroll)
- #centerLayout.addWidget(self.icondock)
- centerLayout.addWidget(iconListContainer,0,Qt.AlignCenter)
- centercontainer = QWidget()
- centercontainer.setLayout(centerLayout)
- # self.scrolldock = QDockWidget('WorkSpace',self)
- # self.scrolldock.setObjectName('WorkSpace')
- # self.scrolldock.setWidget(centercontainer)
- # self.scrolldock.setFeatures(QDockWidget.NoDockWidgetFeatures)
- # orititle = self.scrolldock.titleBarWidget()
- # tmpwidget = QWidget()
- # self.scrolldock.setTitleBarWidget(tmpwidget)
- # del orititle
- self.setCentralWidget(centercontainer) #self.scrolldock
- self.addDockWidget(Qt.RightDockWidgetArea, self.dock)
- # self.filedock.setFeatures(QDockWidget.DockWidgetFloatable)
- self.filedock.setFeatures(self.filedock.features() ^ QDockWidget.DockWidgetFloatable)
- self.dockFeatures = QDockWidget.DockWidgetClosable | QDockWidget.DockWidgetFloatable
- self.dock.setFeatures(self.dock.features() ^ self.dockFeatures)
- self.filedock.setFeatures(QDockWidget.NoDockWidgetFeatures)
- ###### Actions #######
- action = partial(newAction, self)
- quit = action(getStr('quit'), self.close,
- 'Ctrl+Q', 'quit', getStr('quitApp'))
- opendir = action(getStr('openDir'), self.openDirDialog,
- 'Ctrl+u', 'open', getStr('openDir'))
- save = action(getStr('save'), self.saveFile,
- 'Ctrl+V', 'verify', getStr('saveDetail'), enabled=False)
- alcm = action(getStr('choosemodel'), self.autolcm,
- 'Ctrl+M', 'next', getStr('tipchoosemodel'))
- deleteImg = action(getStr('deleteImg'), self.deleteImg, 'Ctrl+Shift+D', 'close', getStr('deleteImgDetail'),
- enabled=True)
- resetAll = action(getStr('resetAll'), self.resetAll, None, 'resetall', getStr('resetAllDetail'))
- color1 = action(getStr('boxLineColor'), self.chooseColor1,
- 'Ctrl+L', 'color_line', getStr('boxLineColorDetail'))
- createMode = action(getStr('crtBox'), self.setCreateMode,
- 'w', 'new', getStr('crtBoxDetail'), enabled=False)
- editMode = action('&Edit\nRectBox', self.setEditMode,
- 'Ctrl+J', 'edit', u'Move and edit Boxs', enabled=False)
- create = action(getStr('crtBox'), self.createShape,
- 'w', 'new', getStr('crtBoxDetail'), enabled=False)
- delete = action(getStr('delBox'), self.deleteSelectedShape,
- 'backspace', 'delete', getStr('delBoxDetail'), enabled=False)
- copy = action(getStr('dupBox'), self.copySelectedShape,
- 'Ctrl+C', 'copy', getStr('dupBoxDetail'),
- enabled=False)
- hideAll = action(getStr('hideBox'), partial(self.togglePolygons, False),
- 'Ctrl+H', 'hide', getStr('hideAllBoxDetail'),
- enabled=False)
- showAll = action(getStr('showBox'), partial(self.togglePolygons, True),
- 'Ctrl+A', 'hide', getStr('showAllBoxDetail'),
- enabled=False)
- help = action(getStr('tutorial'), self.showTutorialDialog, None, 'help', getStr('tutorialDetail'))
- showInfo = action(getStr('info'), self.showInfoDialog, None, 'help', getStr('info'))
- showSteps = action(getStr('steps'), self.showStepsDialog, None, 'help', getStr('steps'))
- zoom = QWidgetAction(self)
- zoom.setDefaultWidget(self.zoomWidget)
- self.zoomWidget.setWhatsThis(
- u"Zoom in or out of the image. Also accessible with"
- " %s and %s from the canvas." % (fmtShortcut("Ctrl+[-+]"),
- fmtShortcut("Ctrl+Wheel")))
- self.zoomWidget.setEnabled(False)
- zoomIn = action(getStr('zoomin'), partial(self.addZoom, 10),
- 'Ctrl++', 'zoom-in', getStr('zoominDetail'), enabled=False)
- zoomOut = action(getStr('zoomout'), partial(self.addZoom, -10),
- 'Ctrl+-', 'zoom-out', getStr('zoomoutDetail'), enabled=False)
- zoomOrg = action(getStr('originalsize'), partial(self.setZoom, 100),
- 'Ctrl+=', 'zoom', getStr('originalsizeDetail'), enabled=False)
- fitWindow = action(getStr('fitWin'), self.setFitWindow,
- 'Ctrl+F', 'fit-window', getStr('fitWinDetail'),
- checkable=True, enabled=False)
- fitWidth = action(getStr('fitWidth'), self.setFitWidth,
- 'Ctrl+Shift+F', 'fit-width', getStr('fitWidthDetail'),
- checkable=True, enabled=False)
- # Group zoom controls into a list for easier toggling.
- zoomActions = (self.zoomWidget, zoomIn, zoomOut,
- zoomOrg, fitWindow, fitWidth)
- self.zoomMode = self.MANUAL_ZOOM
- self.scalers = {
- self.FIT_WINDOW: self.scaleFitWindow,
- self.FIT_WIDTH: self.scaleFitWidth,
- # Set to one to scale to 100% when loading files.
- self.MANUAL_ZOOM: lambda: 1,
- }
- edit = action(getStr('editLabel'), self.editLabel,
- 'Ctrl+E', 'edit', getStr('editLabelDetail'),
- enabled=False)
- ######## New actions #######
- AutoRec = action(getStr('autoRecognition'), self.autoRecognition,
- 'Ctrl+Shift+A', 'Auto', getStr('autoRecognition'), enabled=False)
- reRec = action(getStr('reRecognition'), self.reRecognition,
- 'Ctrl+Shift+R', 'reRec', getStr('reRecognition'), enabled=False)
- singleRere = action(getStr('singleRe'), self.singleRerecognition,
- 'Ctrl+R', 'reRec', getStr('singleRe'), enabled=False)
- createpoly = action(getStr('creatPolygon'), self.createPolygon,
- 'q', 'new', 'Creat Polygon', enabled=True)
- saveRec = action(getStr('saveRec'), self.saveRecResult,
- '', 'save', getStr('saveRec'), enabled=False)
- saveLabel = action(getStr('saveLabel'), self.saveLabelFile, #
- 'Ctrl+S', 'save', getStr('saveLabel'), enabled=False)
- self.editButton.setDefaultAction(edit)
- self.newButton.setDefaultAction(create)
- self.DelButton.setDefaultAction(deleteImg)
- self.SaveButton.setDefaultAction(save)
- self.AutoRecognition.setDefaultAction(AutoRec)
- self.reRecogButton.setDefaultAction(reRec)
- # self.preButton.setDefaultAction(openPrevImg)
- # self.nextButton.setDefaultAction(openNextImg)
- ############# Zoom layout ##############
- zoomLayout = QHBoxLayout()
- zoomLayout.addStretch()
- self.zoominButton = QToolButton()
- self.zoominButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon)
- self.zoominButton.setDefaultAction(zoomIn)
- self.zoomoutButton = QToolButton()
- self.zoomoutButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon)
- self.zoomoutButton.setDefaultAction(zoomOut)
- self.zoomorgButton = QToolButton()
- self.zoomorgButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon)
- self.zoomorgButton.setDefaultAction(zoomOrg)
- zoomLayout.addWidget(self.zoominButton)
- zoomLayout.addWidget(self.zoomorgButton)
- zoomLayout.addWidget(self.zoomoutButton)
- zoomContainer = QWidget()
- zoomContainer.setLayout(zoomLayout)
- zoomContainer.setGeometry(0, 0, 30, 150)
- shapeLineColor = action(getStr('shapeLineColor'), self.chshapeLineColor,
- icon='color_line', tip=getStr('shapeLineColorDetail'),
- enabled=False)
- shapeFillColor = action(getStr('shapeFillColor'), self.chshapeFillColor,
- icon='color', tip=getStr('shapeFillColorDetail'),
- enabled=False)
- # Label list context menu.
- labelMenu = QMenu()
- addActions(labelMenu, (edit, delete))
- self.labelList.setContextMenuPolicy(Qt.CustomContextMenu)
- self.labelList.customContextMenuRequested.connect(
- self.popLabelListMenu)
- # Draw squares/rectangles
- self.drawSquaresOption = QAction(getStr('drawSquares'), self)
- self.drawSquaresOption.setCheckable(True)
- self.drawSquaresOption.setChecked(settings.get(SETTING_DRAW_SQUARE, False))
- self.drawSquaresOption.triggered.connect(self.toogleDrawSquare)
- # Store actions for further handling.
- self.actions = struct(save=save, resetAll=resetAll, deleteImg=deleteImg,
- lineColor=color1, create=create, delete=delete, edit=edit, copy=copy,
- saveRec=saveRec, singleRere=singleRere,AutoRec=AutoRec,reRec=reRec,
- createMode=createMode, editMode=editMode,
- shapeLineColor=shapeLineColor, shapeFillColor=shapeFillColor,
- zoom=zoom, zoomIn=zoomIn, zoomOut=zoomOut, zoomOrg=zoomOrg,
- fitWindow=fitWindow, fitWidth=fitWidth,
- zoomActions=zoomActions, saveLabel=saveLabel,
- fileMenuActions=(
- opendir, saveLabel, resetAll, quit),
- beginner=(), advanced=(),
- editMenu=(createpoly, edit, copy, delete,singleRere,
- None, color1, self.drawSquaresOption),
- beginnerContext=(create, edit, copy, delete, singleRere),
- advancedContext=(createMode, editMode, edit, copy,
- delete, shapeLineColor, shapeFillColor),
- onLoadActive=(
- create, createMode, editMode),
- onShapesPresent=(hideAll, showAll))
- # menus
- self.menus = struct(
- file=self.menu('&'+getStr('mfile')),
- edit=self.menu('&'+getStr('medit')),
- view=self.menu('&'+getStr('mview')),
- autolabel=self.menu('&PaddleOCR'),
- help=self.menu('&'+getStr('mhelp')),
- recentFiles=QMenu('Open &Recent'),
- labelList=labelMenu)
- self.lastLabel = None
- # Add option to enable/disable labels being displayed at the top of bounding boxes
- self.displayLabelOption = QAction(getStr('displayLabel'), self)
- self.displayLabelOption.setShortcut("Ctrl+Shift+P")
- self.displayLabelOption.setCheckable(True)
- self.displayLabelOption.setChecked(settings.get(SETTING_PAINT_LABEL, False))
- self.displayLabelOption.triggered.connect(self.togglePaintLabelsOption)
- self.labelDialogOption = QAction(getStr('labelDialogOption'), self)
- self.labelDialogOption.setShortcut("Ctrl+Shift+L")
- self.labelDialogOption.setCheckable(True)
- self.labelDialogOption.setChecked(settings.get(SETTING_PAINT_LABEL, False))
- self.labelDialogOption.triggered.connect(self.speedChoose)
- addActions(self.menus.file,
- (opendir, None, saveLabel, saveRec, None, resetAll, deleteImg, quit))
- addActions(self.menus.help, (showSteps, showInfo))
- addActions(self.menus.view, (
- self.displayLabelOption, self.labelDialogOption,
- None,
- hideAll, showAll, None,
- zoomIn, zoomOut, zoomOrg, None,
- fitWindow, fitWidth))
- addActions(self.menus.autolabel, (AutoRec, reRec, alcm, None, help)) #
- self.menus.file.aboutToShow.connect(self.updateFileMenu)
- # Custom context menu for the canvas widget:
- addActions(self.canvas.menus[0], self.actions.beginnerContext)
- addActions(self.canvas.menus[1], (
- action('&Copy here', self.copyShape),
- action('&Move here', self.moveShape)))
- self.statusBar().showMessage('%s started.' % __appname__)
- self.statusBar().show()
- # Application state.
- self.image = QImage()
- self.filePath = ustr(defaultFilename)
- self.lastOpenDir = None
- self.recentFiles = []
- self.maxRecent = 7
- self.lineColor = None
- self.fillColor = None
- self.zoom_level = 100
- self.fit_window = False
- # Add Chris
- self.difficult = False
- ## Fix the compatible issue for qt4 and qt5. Convert the QStringList to python list
- if settings.get(SETTING_RECENT_FILES):
- if have_qstring():
- recentFileQStringList = settings.get(SETTING_RECENT_FILES)
- self.recentFiles = [ustr(i) for i in recentFileQStringList]
- else:
- self.recentFiles = recentFileQStringList = settings.get(SETTING_RECENT_FILES)
- size = settings.get(SETTING_WIN_SIZE, QSize(1200, 800))
- position = QPoint(0, 0)
- saved_position = settings.get(SETTING_WIN_POSE, position)
- # Fix the multiple monitors issue
- for i in range(QApplication.desktop().screenCount()):
- if QApplication.desktop().availableGeometry(i).contains(saved_position):
- position = saved_position
- break
- self.resize(size)
- self.move(position)
- saveDir = ustr(settings.get(SETTING_SAVE_DIR, None))
- self.lastOpenDir = ustr(settings.get(SETTING_LAST_OPEN_DIR, None))
- self.restoreState(settings.get(SETTING_WIN_STATE, QByteArray()))
- Shape.line_color = self.lineColor = QColor(settings.get(SETTING_LINE_COLOR, DEFAULT_LINE_COLOR))
- Shape.fill_color = self.fillColor = QColor(settings.get(SETTING_FILL_COLOR, DEFAULT_FILL_COLOR))
- self.canvas.setDrawingColor(self.lineColor)
- # Add chris
- Shape.difficult = self.difficult
- # ADD:
- # Populate the File menu dynamically.
- self.updateFileMenu()
- # Since loading the file may take some time, make sure it runs in the background.
- if self.filePath and os.path.isdir(self.filePath):
- self.queueEvent(partial(self.importDirImages, self.filePath or ""))
- elif self.filePath:
- self.queueEvent(partial(self.loadFile, self.filePath or ""))
- # Callbacks:
- self.zoomWidget.valueChanged.connect(self.paintCanvas)
- self.populateModeActions()
- # Display cursor coordinates at the right of status bar
- self.labelCoordinates = QLabel('')
- self.statusBar().addPermanentWidget(self.labelCoordinates)
- # Open Dir if deafult file
- if self.filePath and os.path.isdir(self.filePath):
- self.openDirDialog(dirpath=self.filePath, silent=True)
- def keyReleaseEvent(self, event):
- if event.key() == Qt.Key_Control:
- self.canvas.setDrawingShapeToSquare(False)
- def keyPressEvent(self, event):
- if event.key() == Qt.Key_Control:
- # Draw rectangle if Ctrl is pressed
- self.canvas.setDrawingShapeToSquare(True)
- def noShapes(self):
- return not self.itemsToShapes
- def populateModeActions(self):
- self.canvas.menus[0].clear()
- addActions(self.canvas.menus[0], self.actions.beginnerContext)
- self.menus.edit.clear()
- actions = (self.actions.create,) # if self.beginner() else (self.actions.createMode, self.actions.editMode)
- addActions(self.menus.edit, actions + self.actions.editMenu)
- def setDirty(self):
- self.dirty = True
- self.actions.save.setEnabled(True)
- def setClean(self):
- self.dirty = False
- self.actions.save.setEnabled(False)
- self.actions.create.setEnabled(True)
- def toggleActions(self, value=True):
- """Enable/Disable widgets which depend on an opened image."""
- for z in self.actions.zoomActions:
- z.setEnabled(value)
- for action in self.actions.onLoadActive:
- action.setEnabled(value)
- def queueEvent(self, function):
- QTimer.singleShot(0, function)
- def status(self, message, delay=5000):
- self.statusBar().showMessage(message, delay)
- def resetState(self):
- self.itemsToShapes.clear()
- self.shapesToItems.clear()
- self.itemsToShapesbox.clear() # ADD
- self.shapesToItemsbox.clear()
- self.labelList.clear()
- self.BoxList.clear()
- self.filePath = None
- self.imageData = None
- self.labelFile = None
- self.canvas.resetState()
- self.labelCoordinates.clear()
- # self.comboBox.cb.clear()
- self.result_dic = []
- def currentItem(self):
- items = self.labelList.selectedItems()
- if items:
- return items[0]
- return None
- def currentBox(self):
- items = self.BoxList.selectedItems()
- if items:
- return items[0]
- return None
- def addRecentFile(self, filePath):
- if filePath in self.recentFiles:
- self.recentFiles.remove(filePath)
- elif len(self.recentFiles) >= self.maxRecent:
- self.recentFiles.pop()
- self.recentFiles.insert(0, filePath)
- def beginner(self):
- return self._beginner
- def advanced(self):
- return not self.beginner()
- def getAvailableScreencastViewer(self):
- osName = platform.system()
- if osName == 'Windows':
- return ['C:\\Program Files\\Internet Explorer\\iexplore.exe']
- elif osName == 'Linux':
- return ['xdg-open']
- elif osName == 'Darwin':
- return ['open']
- ## Callbacks ##
- def showTutorialDialog(self):
- subprocess.Popen(self.screencastViewer + [self.screencast])
- def showInfoDialog(self):
- from libs.__init__ import __version__
- msg = u'Name:{0} \nApp Version:{1} \n{2} '.format(__appname__, __version__, sys.version_info)
- QMessageBox.information(self, u'Information', msg)
- def showStepsDialog(self):
- msg = stepsInfo(self.lang)
- QMessageBox.information(self, u'Information', msg)
- def createShape(self):
- assert self.beginner()
- self.canvas.setEditing(False)
- self.actions.create.setEnabled(False)
- self.canvas.fourpoint = False
- def createPolygon(self):
- assert self.beginner()
- self.canvas.setEditing(False)
- self.canvas.fourpoint = True
- self.actions.create.setEnabled(False)
- def toggleDrawingSensitive(self, drawing=True):
- """In the middle of drawing, toggling between modes should be disabled."""
- self.actions.editMode.setEnabled(not drawing)
- if not drawing and self.beginner():
- # Cancel creation.
- print('Cancel creation.')
- self.canvas.setEditing(True)
- self.canvas.restoreCursor()
- self.actions.create.setEnabled(True)
- def toggleDrawMode(self, edit=True):
- self.canvas.setEditing(edit)
- self.actions.createMode.setEnabled(edit)
- self.actions.editMode.setEnabled(not edit)
- def setCreateMode(self):
- assert self.advanced()
- self.toggleDrawMode(False)
- def setEditMode(self):
- assert self.advanced()
- self.toggleDrawMode(True)
- self.labelSelectionChanged()
- def updateFileMenu(self):
- currFilePath = self.filePath
- def exists(filename):
- return os.path.exists(filename)
- menu = self.menus.recentFiles
- menu.clear()
- files = [f for f in self.recentFiles if f !=
- currFilePath and exists(f)]
- for i, f in enumerate(files):
- icon = newIcon('labels')
- action = QAction(
- icon, '&%d %s' % (i + 1, QFileInfo(f).fileName()), self)
- action.triggered.connect(partial(self.loadRecent, f))
- menu.addAction(action)
- def popLabelListMenu(self, point):
- self.menus.labelList.exec_(self.labelList.mapToGlobal(point))
- def editLabel(self):
- if not self.canvas.editing():
- return
- item = self.currentItem()
- if not item:
- return
- text = self.labelDialog.popUp(item.text())
- if text is not None:
- item.setText(text)
- # item.setBackground(generateColorByText(text))
- self.setDirty()
- self.updateComboBox()
- ######## detection box related functions #######
- def boxItemChanged(self, item):
- shape = self.itemsToShapesbox[item]
- box = ast.literal_eval(item.text())
- # print('shape in labelItemChanged is',shape.points)
- if box != [(p.x(), p.y()) for p in shape.points]:
- # shape.points = box
- shape.points = [QPointF(p[0], p[1]) for p in box]
- # QPointF(x,y)
- # shape.line_color = generateColorByText(shape.label)
- self.setDirty()
- else: # User probably changed item visibility
- self.canvas.setShapeVisible(shape, True)#item.checkState() == Qt.Checked
- def editBox(self): # ADD
- if not self.canvas.editing():
- return
- item = self.currentBox()
- if not item:
- return
- text = self.labelDialog.popUp(item.text())
- imageSize = str(self.image.size())
- width, height = self.image.width(), self.image.height()
- if text:
- try:
- text_list = eval(text)
- except:
- msg_box = QMessageBox(QMessageBox.Warning, 'Warning', 'Please enter the correct format')
- msg_box.exec_()
- return
- if len(text_list) < 4:
- msg_box = QMessageBox(QMessageBox.Warning, 'Warning', 'Please enter the coordinates of 4 points')
- msg_box.exec_()
- return
- for box in text_list:
- if box[0] > width or box[0] < 0 or box[1] > height or box[1] < 0:
- msg_box = QMessageBox(QMessageBox.Warning, 'Warning', 'Out of picture size')
- msg_box.exec_()
- return
- item.setText(text)
- # item.setBackground(generateColorByText(text))
- self.setDirty()
- self.updateComboBox()
- def updateBoxlist(self):
- shape = self.canvas.selectedShape
- item = self.shapesToItemsbox[shape] # listitem
- text = [(int(p.x()), int(p.y())) for p in shape.points]
- item.setText(str(text))
- self.setDirty()
- def indexTo5Files(self, currIndex):
- if currIndex < 2:
- return self.mImgList[:5]
- elif currIndex > len(self.mImgList)-3:
- return self.mImgList[-5:]
- else:
- return self.mImgList[currIndex - 2 : currIndex + 3]
- # Tzutalin 20160906 : Add file list and dock to move faster
- def fileitemDoubleClicked(self, item=None):
- self.currIndex = self.mImgList.index(ustr(os.path.join(os.path.abspath(self.dirname), item.text())))
- filename = self.mImgList[self.currIndex]
- if filename:
- self.mImgList5 = self.indexTo5Files(self.currIndex)
- # self.additems5(None)
- self.loadFile(filename)
- def iconitemDoubleClicked(self, item=None):
- self.currIndex = self.mImgList.index(ustr(os.path.join(item.toolTip())))
- filename = self.mImgList[self.currIndex]
- if filename:
- self.mImgList5 = self.indexTo5Files(self.currIndex)
- # self.additems5(None)
- self.loadFile(filename)
- def CanvasSizeChange(self):
- if len(self.mImgList) > 0:
- self.zoomWidget.setValue(self.zoomWidgetValue + self.imgsplider.value())
- # React to canvas signals.
- def shapeSelectionChanged(self, selected=False):
- if self._noSelectionSlot:
- self._noSelectionSlot = False
- else:
- shape = self.canvas.selectedShape
- if shape:
- self.shapesToItems[shape].setSelected(True)
- self.shapesToItemsbox[shape].setSelected(True) # ADD
- else:
- self.labelList.clearSelection()
- self.actions.delete.setEnabled(selected)
- self.actions.copy.setEnabled(selected)
- self.actions.edit.setEnabled(selected)
- self.actions.shapeLineColor.setEnabled(selected)
- self.actions.shapeFillColor.setEnabled(selected)
- self.actions.singleRere.setEnabled(selected)
- def addLabel(self, shape):
- shape.paintLabel = self.displayLabelOption.isChecked()
- item = HashableQListWidgetItem(shape.label)
- item.setFlags(item.flags() | Qt.ItemIsUserCheckable)
- item.setCheckState(Qt.Unchecked) if shape.difficult else item.setCheckState(Qt.Checked)
- # Checked means difficult is False
- # item.setBackground(generateColorByText(shape.label))
- self.itemsToShapes[item] = shape
- self.shapesToItems[shape] = item
- self.labelList.addItem(item)
- # print('item in add label is ',[(p.x(), p.y()) for p in shape.points], shape.label)
- # ADD for box
- item = HashableQListWidgetItem(str([(int(p.x()), int(p.y())) for p in shape.points]))
- self.itemsToShapesbox[item] = shape
- self.shapesToItemsbox[shape] = item
- self.BoxList.addItem(item)
- for action in self.actions.onShapesPresent:
- action.setEnabled(True)
- self.updateComboBox()
- def remLabel(self, shape):
- if shape is None:
- # print('rm empty label')
- return
- item = self.shapesToItems[shape]
- self.labelList.takeItem(self.labelList.row(item))
- del self.shapesToItems[shape]
- del self.itemsToShapes[item]
- self.updateComboBox()
- # ADD:
- item = self.shapesToItemsbox[shape]
- self.BoxList.takeItem(self.BoxList.row(item))
- del self.shapesToItemsbox[shape]
- del self.itemsToShapesbox[item]
- self.updateComboBox()
- def loadLabels(self, shapes):
- s = []
- for label, points, line_color, fill_color, difficult in shapes:
- shape = Shape(label=label)
- for x, y in points:
- # Ensure the labels are within the bounds of the image. If not, fix them.
- x, y, snapped = self.canvas.snapPointToCanvas(x, y)
- if snapped:
- self.setDirty()
- shape.addPoint(QPointF(x, y))
- shape.difficult = difficult
- shape.close()
- s.append(shape)
- # if line_color:
- # shape.line_color = QColor(*line_color)
- # else:
- # shape.line_color = generateColorByText(label)
- #
- # if fill_color:
- # shape.fill_color = QColor(*fill_color)
- # else:
- # shape.fill_color = generateColorByText(label)
- self.addLabel(shape)
- self.updateComboBox()
- self.canvas.loadShapes(s)
- def singleLabel(self, shape):
- if shape is None:
- # print('rm empty label')
- return
- item = self.shapesToItems[shape]
- item.setText(shape.label)
- self.updateComboBox()
- # ADD:
- item = self.shapesToItemsbox[shape]
- item.setText(str([(int(p.x()), int(p.y())) for p in shape.points]))
- self.updateComboBox()
- def updateComboBox(self):
- # Get the unique labels and add them to the Combobox.
- itemsTextList = [str(self.labelList.item(i).text()) for i in range(self.labelList.count())]
- uniqueTextList = list(set(itemsTextList))
- # Add a null row for showing all the labels
- uniqueTextList.append("")
- uniqueTextList.sort()
- # self.comboBox.update_items(uniqueTextList)
- def saveLabels(self, annotationFilePath, mode='Auto'):
- # Mode is Auto means that labels will be loaded from self.result_dic totally, which is the output of ocr model
- annotationFilePath = ustr(annotationFilePath)
- def format_shape(s):
- # print('s in saveLabels is ',s)
- return dict(label=s.label, # str
- line_color=s.line_color.getRgb(),
- fill_color=s.fill_color.getRgb(),
- points=[(p.x(), p.y()) for p in s.points], # QPonitF
- # add chris
- difficult=s.difficult) # bool
- shapes = [] if mode == 'Auto' else \
- [format_shape(shape) for shape in self.canvas.shapes]
- # Can add differrent annotation formats here
- for box in self.result_dic:
- trans_dic = {"label": box[1][0], "points": box[0], 'difficult': False}
- if trans_dic["label"] == "" and mode == 'Auto':
- continue
- shapes.append(trans_dic)
- try:
- trans_dic = []
- for box in shapes:
- trans_dic.append({"transcription": box['label'], "points": box['points'], 'difficult': box['difficult']})
- self.PPlabel[annotationFilePath] = trans_dic
- if mode == 'Auto':
- self.Cachelabel[annotationFilePath] = trans_dic
- # else:
- # self.labelFile.save(annotationFilePath, shapes, self.filePath, self.imageData,
- # self.lineColor.getRgb(), self.fillColor.getRgb())
- # print('Image:{0} -> Annotation:{1}'.format(self.filePath, annotationFilePath))
- return True
- except:
- self.errorMessage(u'Error saving label data')
- return False
- def copySelectedShape(self):
- self.addLabel(self.canvas.copySelectedShape())
- # fix copy and delete
- self.shapeSelectionChanged(True)
- def labelSelectionChanged(self):
- item = self.currentItem()
- self.labelList.scrollToItem(item, QAbstractItemView.EnsureVisible)
- if item and self.canvas.editing():
- self._noSelectionSlot = True
- self.canvas.selectShape(self.itemsToShapes[item])
- shape = self.itemsToShapes[item]
- def boxSelectionChanged(self):
- item = self.currentBox()
- self.BoxList.scrollToItem(item, QAbstractItemView.EnsureVisible)
- if item and self.canvas.editing():
- self._noSelectionSlot = True
- self.canvas.selectShape(self.itemsToShapesbox[item])
- shape = self.itemsToShapesbox[item]
- def labelItemChanged(self, item):
- shape = self.itemsToShapes[item]
- label = item.text()
- if label != shape.label:
- shape.label = item.text()
- # shape.line_color = generateColorByText(shape.label)
- self.setDirty()
- elif not ((item.checkState()== Qt.Unchecked) ^ (not shape.difficult)):
- shape.difficult = True if item.checkState() == Qt.Unchecked else False
- self.setDirty()
- else: # User probably changed item visibility
- self.canvas.setShapeVisible(shape, True) # item.checkState() == Qt.Checked
- # self.actions.save.setEnabled(True)
- # Callback functions:
- def newShape(self, value=True):
- """Pop-up and give focus to the label editor.
- position MUST be in global coordinates.
- """
- if len(self.labelHist) > 0:
- self.labelDialog = LabelDialog(
- parent=self, listItem=self.labelHist)
- if value:
- text = self.labelDialog.popUp(text=self.prevLabelText)
- self.lastLabel = text
- else:
- text = self.prevLabelText
- if text is not None:
- self.prevLabelText = self.stringBundle.getString('tempLabel')
- # generate_color = generateColorByText(text)
- shape = self.canvas.setLastLabel(text, None, None)#generate_color, generate_color
- self.addLabel(shape)
- if self.beginner(): # Switch to edit mode.
- self.canvas.setEditing(True)
- self.actions.create.setEnabled(True)
- else:
- self.actions.editMode.setEnabled(True)
- self.setDirty()
- else:
- # self.canvas.undoLastLine()
- self.canvas.resetAllLines()
- def scrollRequest(self, delta, orientation):
- units = - delta / (8 * 15)
- bar = self.scrollBars[orientation]
- bar.setValue(bar.value() + bar.singleStep() * units)
- def setZoom(self, value):
- self.actions.fitWidth.setChecked(False)
- self.actions.fitWindow.setChecked(False)
- self.zoomMode = self.MANUAL_ZOOM
- self.zoomWidget.setValue(value)
- def addZoom(self, increment=10):
- self.setZoom(self.zoomWidget.value() + increment)
- def zoomRequest(self, delta):
- # get the current scrollbar positions
- # calculate the percentages ~ coordinates
- h_bar = self.scrollBars[Qt.Horizontal]
- v_bar = self.scrollBars[Qt.Vertical]
- # get the current maximum, to know the difference after zooming
- h_bar_max = h_bar.maximum()
- v_bar_max = v_bar.maximum()
- # get the cursor position and canvas size
- # calculate the desired movement from 0 to 1
- # where 0 = move left
- # 1 = move right
- # up and down analogous
- cursor = QCursor()
- pos = cursor.pos()
- relative_pos = QWidget.mapFromGlobal(self, pos)
- cursor_x = relative_pos.x()
- cursor_y = relative_pos.y()
- w = self.scrollArea.width()
- h = self.scrollArea.height()
- # the scaling from 0 to 1 has some padding
- # you don't have to hit the very leftmost pixel for a maximum-left movement
- margin = 0.1
- move_x = (cursor_x - margin * w) / (w - 2 * margin * w)
- move_y = (cursor_y - margin * h) / (h - 2 * margin * h)
- # clamp the values from 0 to 1
- move_x = min(max(move_x, 0), 1)
- move_y = min(max(move_y, 0), 1)
- # zoom in
- units = delta / (8 * 15)
- scale = 10
- self.addZoom(scale * units)
- # get the difference in scrollbar values
- # this is how far we can move
- d_h_bar_max = h_bar.maximum() - h_bar_max
- d_v_bar_max = v_bar.maximum() - v_bar_max
- # get the new scrollbar values
- new_h_bar_value = h_bar.value() + move_x * d_h_bar_max
- new_v_bar_value = v_bar.value() + move_y * d_v_bar_max
- h_bar.setValue(new_h_bar_value)
- v_bar.setValue(new_v_bar_value)
- def setFitWindow(self, value=True):
- if value:
- self.actions.fitWidth.setChecked(False)
- self.zoomMode = self.FIT_WINDOW if value else self.MANUAL_ZOOM
- self.adjustScale()
- def setFitWidth(self, value=True):
- if value:
- self.actions.fitWindow.setChecked(False)
- self.zoomMode = self.FIT_WIDTH if value else self.MANUAL_ZOOM
- self.adjustScale()
- def togglePolygons(self, value):
- for item, shape in self.itemsToShapes.items():
- self.canvas.setShapeVisible(shape, value)
- def loadFile(self, filePath=None):
- """Load the specified file, or the last opened file if None."""
- self.resetState()
- self.canvas.setEnabled(False)
- if filePath is None:
- filePath = self.settings.get(SETTING_FILENAME)
- # Make sure that filePath is a regular python string, rather than QString
- filePath = ustr(filePath)
- # Fix bug: An index error after select a directory when open a new file.
- unicodeFilePath = ustr(filePath)
- # unicodeFilePath = os.path.abspath(unicodeFilePath)
- # Tzutalin 20160906 : Add file list and dock to move faster
- # Highlight the file item
- if unicodeFilePath and self.fileListWidget.count() > 0:
- if unicodeFilePath in self.mImgList:
- index = self.mImgList.index(unicodeFilePath)
- fileWidgetItem = self.fileListWidget.item(index)
- print('unicodeFilePath is', unicodeFilePath)
- fileWidgetItem.setSelected(True)
- ###
- self.iconlist.clear()
- self.additems5(None)
- for i in range(5):
- item_tooltip = self.iconlist.item(i).toolTip()
- # print(i,"---",item_tooltip)
- if item_tooltip == ustr(filePath):
- titem = self.iconlist.item(i)
- titem.setSelected(True)
- self.iconlist.scrollToItem(titem)
- break
- else:
- self.fileListWidget.clear()
- self.mImgList.clear()
- self.iconlist.clear()
- # if unicodeFilePath and self.iconList.count() > 0:
- # if unicodeFilePath in self.mImgList:
-
- if unicodeFilePath and os.path.exists(unicodeFilePath):
- self.canvas.verified = False
- cvimg = cv2.imdecode(np.fromfile(unicodeFilePath, dtype=np.uint8), 1)
- height, width, depth = cvimg.shape
- cvimg = cv2.cvtColor(cvimg, cv2.COLOR_BGR2RGB)
- image = QImage(cvimg.data, width, height, width * depth, QImage.Format_RGB888)
- if image.isNull():
- self.errorMessage(u'Error opening file',
- u"<p>Make sure <i>%s</i> is a valid image file." % unicodeFilePath)
- self.status("Error reading %s" % unicodeFilePath)
- return False
- self.status("Loaded %s" % os.path.basename(unicodeFilePath))
- self.image = image
- self.filePath = unicodeFilePath
- self.canvas.loadPixmap(QPixmap.fromImage(image))
- if self.validFilestate(filePath) is True:
- self.setClean()
- else:
- self.dirty = False
- self.actions.save.setEnabled(True)
- self.canvas.setEnabled(True)
- self.adjustScale(initial=True)
- self.paintCanvas()
- self.addRecentFile(self.filePath)
- self.toggleActions(True)
- self.showBoundingBoxFromPPlabel(filePath)
-
- self.setWindowTitle(__appname__ + ' ' + filePath)
- # Default : select last item if there is at least one item
- if self.labelList.count():
- self.labelList.setCurrentItem(self.labelList.item(self.labelList.count() - 1))
- self.labelList.item(self.labelList.count() - 1).setSelected(True)
- self.canvas.setFocus(True)
- return True
- return False
-
- def showBoundingBoxFromPPlabel(self, filePath):
- imgidx = self.getImglabelidx(filePath)
- if imgidx not in self.PPlabel.keys():
- return
- shapes = []
- for box in self.PPlabel[imgidx]:
- shapes.append((box['transcription'], box['points'], None, None, box['difficult']))
- print(shapes)
- self.loadLabels(shapes)
- self.canvas.verified = False
- def validFilestate(self, filePath):
- if filePath not in self.fileStatedict.keys():
- return None
- elif self.fileStatedict[filePath] == 1:
- return True
- else:
- return False
- def resizeEvent(self, event):
- if self.canvas and not self.image.isNull() \
- and self.zoomMode != self.MANUAL_ZOOM:
- self.adjustScale()
- super(MainWindow, self).resizeEvent(event)
- def paintCanvas(self):
- assert not self.image.isNull(), "cannot paint null image"
- self.canvas.scale = 0.01 * self.zoomWidget.value()
- self.canvas.adjustSize()
- self.canvas.update()
- def adjustScale(self, initial=False):
- value = self.scalers[self.FIT_WINDOW if initial else self.zoomMode]()
- self.zoomWidget.setValue(int(100 * value))
- def scaleFitWindow(self):
- """Figure out the size of the pixmap in order to fit the main widget."""
- e = 2.0 # So that no scrollbars are generated.
- w1 = self.centralWidget().width() - e
- h1 = self.centralWidget().height() - e -110
- a1 = w1 / h1
- # Calculate a new scale value based on the pixmap's aspect ratio.
- w2 = self.canvas.pixmap.width() - 0.0
- h2 = self.canvas.pixmap.height() - 0.0
- a2 = w2 / h2
- return w1 / w2 if a2 >= a1 else h1 / h2
- def scaleFitWidth(self):
- # The epsilon does not seem to work too well here.
- w = self.centralWidget().width() - 2.0
- return w / self.canvas.pixmap.width()
- def closeEvent(self, event):
- if not self.mayContinue():
- event.ignore()
- else:
- settings = self.settings
- # If it loads images from dir, don't load it at the begining
- if self.dirname is None:
- settings[SETTING_FILENAME] = self.filePath if self.filePath else ''
- else:
- settings[SETTING_FILENAME] = ''
- settings[SETTING_WIN_SIZE] = self.size()
- settings[SETTING_WIN_POSE] = self.pos()
- settings[SETTING_WIN_STATE] = self.saveState()
- settings[SETTING_LINE_COLOR] = self.lineColor
- settings[SETTING_FILL_COLOR] = self.fillColor
- settings[SETTING_RECENT_FILES] = self.recentFiles
- settings[SETTING_ADVANCE_MODE] = not self._beginner
- if self.defaultSaveDir and os.path.exists(self.defaultSaveDir):
- settings[SETTING_SAVE_DIR] = ustr(self.defaultSaveDir)
- else:
- settings[SETTING_SAVE_DIR] = ''
- if self.lastOpenDir and os.path.exists(self.lastOpenDir):
- settings[SETTING_LAST_OPEN_DIR] = self.lastOpenDir
- else:
- settings[SETTING_LAST_OPEN_DIR] = ''
- settings[SETTING_PAINT_LABEL] = self.displayLabelOption.isChecked()
- settings[SETTING_DRAW_SQUARE] = self.drawSquaresOption.isChecked()
- settings.save()
- try:
- self.saveLabelFile()
- except:
- pass
- def loadRecent(self, filename):
- if self.mayContinue():
- self.loadFile(filename)
- def scanAllImages(self, folderPath):
- extensions = ['.%s' % fmt.data().decode("ascii").lower() for fmt in QImageReader.supportedImageFormats()]
- images = []
- for file in os.listdir(folderPath):
- if file.lower().endswith(tuple(extensions)):
- relativePath = os.path.join(folderPath, file)
- path = ustr(os.path.abspath(relativePath))
- images.append(path)
- natural_sort(images, key=lambda x: x.lower())
- return images
- def openDirDialog(self, _value=False, dirpath=None, silent=False):
- if not self.mayContinue():
- return
- defaultOpenDirPath = dirpath if dirpath else '.'
- if self.lastOpenDir and os.path.exists(self.lastOpenDir):
- defaultOpenDirPath = self.lastOpenDir
- else:
- defaultOpenDirPath = os.path.dirname(self.filePath) if self.filePath else '.'
- if silent != True:
- targetDirPath = ustr(QFileDialog.getExistingDirectory(self,
- '%s - Open Directory' % __appname__,
- defaultOpenDirPath,
- QFileDialog.ShowDirsOnly | QFileDialog.DontResolveSymlinks))
- else:
- targetDirPath = ustr(defaultOpenDirPath)
- self.lastOpenDir = targetDirPath
- self.importDirImages(targetDirPath)
- def importDirImages(self, dirpath, isDelete = False):
- if not self.mayContinue() or not dirpath:
- return
- if self.defaultSaveDir and self.defaultSaveDir != dirpath:
- self.saveLabelFile()
- if not isDelete:
- self.loadFilestate(dirpath)
- self.PPlabelpath = dirpath+ '/Label.txt'
- self.PPlabel = self.loadLabelFile(self.PPlabelpath)
- self.Cachelabelpath = dirpath + '/Cache.cach'
- self.Cachelabel = self.loadLabelFile(self.Cachelabelpath)
- if self.Cachelabel:
- self.PPlabel = dict(self.Cachelabel, **self.PPlabel)
- self.lastOpenDir = dirpath
- self.dirname = dirpath
- self.defaultSaveDir = dirpath
- self.statusBar().showMessage('%s started. Annotation will be saved to %s' %
- (__appname__, self.defaultSaveDir))
- self.statusBar().show()
- self.filePath = None
- self.fileListWidget.clear()
- self.mImgList = self.scanAllImages(dirpath)
- self.mImgList5 = self.mImgList[:5]
- self.openNextImg()
- doneicon = newIcon('done')
- closeicon = newIcon('close')
- for imgPath in self.mImgList:
- filename = os.path.basename(imgPath)
- if self.validFilestate(imgPath) is True:
- item = QListWidgetItem(doneicon, filename)
- else:
- item = QListWidgetItem(closeicon, filename)
- self.fileListWidget.addItem(item)
- print('DirPath in importDirImages is', dirpath)
- self.iconlist.clear()
- self.additems5(dirpath)
- self.changeFileFolder = True
- self.haveAutoReced = False
- self.AutoRecognition.setEnabled(True)
- self.reRecogButton.setEnabled(True)
- self.actions.AutoRec.setEnabled(True)
- self.actions.reRec.setEnabled(True)
- def openPrevImg(self, _value=False):
- if len(self.mImgList) <= 0:
- return
- if self.filePath is None:
- return
-
- currIndex = self.mImgList.index(self.filePath)
- self.mImgList5 = self.mImgList[:5]
- if currIndex - 1 >= 0:
- filename = self.mImgList[currIndex - 1]
- self.mImgList5 = self.indexTo5Files(currIndex - 1)
- if filename:
- self.loadFile(filename)
- def openNextImg(self, _value=False):
- if not self.mayContinue():
- return
- if len(self.mImgList) <= 0:
- return
- filename = None
- if self.filePath is None:
- filename = self.mImgList[0]
- self.mImgList5 = self.mImgList[:5]
- else:
- currIndex = self.mImgList.index(self.filePath)
- if currIndex + 1 < len(self.mImgList):
- filename = self.mImgList[currIndex + 1]
- self.mImgList5 = self.indexTo5Files(currIndex + 1)
- else:
- self.mImgList5 = self.indexTo5Files(currIndex)
- if filename:
- print('file name in openNext is ',filename)
- self.loadFile(filename)
-
- def updateFileListIcon(self, filename):
- pass
- def saveFile(self, _value=False, mode='Manual'):
- # Manual mode is used for users click "Save" manually,which will change the state of the image
- if self.filePath:
- imgidx = self.getImglabelidx(self.filePath)
- self._saveFile(imgidx, mode=mode)
- def saveFileAs(self, _value=False):
- assert not self.image.isNull(), "cannot save empty image"
- self._saveFile(self.saveFileDialog())
- def saveFileDialog(self, removeExt=True):
- caption = '%s - Choose File' % __appname__
- filters = 'File (*%s)' % LabelFile.suffix
- openDialogPath = self.currentPath()
- dlg = QFileDialog(self, caption, openDialogPath, filters)
- dlg.setDefaultSuffix(LabelFile.suffix[1:])
- dlg.setAcceptMode(QFileDialog.AcceptSave)
- filenameWithoutExtension = os.path.splitext(self.filePath)[0]
- dlg.selectFile(filenameWithoutExtension)
- dlg.setOption(QFileDialog.DontUseNativeDialog, False)
- if dlg.exec_():
- fullFilePath = ustr(dlg.selectedFiles()[0])
- if removeExt:
- return os.path.splitext(fullFilePath)[0] # Return file path without the extension.
- else:
- return fullFilePath
- return ''
- def _saveFile(self, annotationFilePath, mode='Manual'):
- if mode == 'Manual':
- if annotationFilePath and self.saveLabels(annotationFilePath, mode=mode):
- self.setClean()
- self.statusBar().showMessage('Saved to %s' % annotationFilePath)
- self.statusBar().show()
- currIndex = self.mImgList.index(self.filePath)
- item = self.fileListWidget.item(currIndex)
- item.setIcon(newIcon('done'))
- self.fileStatedict[self.filePath] = 1
- if len(self.fileStatedict)%self.autoSaveNum ==0:
- self.saveFilestate()
- self.savePPlabel(mode='Auto')
- self.fileListWidget.insertItem(int(currIndex), item)
- self.openNextImg()
- self.actions.saveRec.setEnabled(True)
- elif mode == 'Auto':
- if annotationFilePath and self.saveLabels(annotationFilePath, mode=mode):
- self.setClean()
- self.statusBar().showMessage('Saved to %s' % annotationFilePath)
- self.statusBar().show()
- def closeFile(self, _value=False):
- if not self.mayContinue():
- return
- self.resetState()
- self.setClean()
- self.toggleActions(False)
- self.canvas.setEnabled(False)
- self.actions.saveAs.setEnabled(False)
- def deleteImg(self):
- deletePath = self.filePath
- if deletePath is not None:
- deleteInfo = self.deleteImgDialog()
- if deleteInfo == QMessageBox.Yes:
- if platform.system() == 'Windows':
- from win32com.shell import shell, shellcon
- shell.SHFileOperation((0, shellcon.FO_DELETE, deletePath, None,
- shellcon.FOF_SILENT | shellcon.FOF_ALLOWUNDO | shellcon.FOF_NOCONFIRMATION,
- None, None))
- # linux
- elif platform.system() == 'Linux':
- cmd = 'trash ' + deletePath
- os.system(cmd)
- # macOS
- elif platform.system() == 'Darwin':
- import subprocess
- absPath = os.path.abspath(deletePath).replace('\\', '\\\\').replace('"', '\\"')
- cmd = ['osascript', '-e',
- 'tell app "Finder" to move {the POSIX file "' + absPath + '"} to trash']
- print(cmd)
- subprocess.call(cmd, stdout=open(os.devnull, 'w'))
- if self.filePath in self.fileStatedict.keys():
- self.fileStatedict.pop(self.filePath)
- imgidx = self.getImglabelidx(self.filePath)
- if imgidx in self.PPlabel.keys():
- self.PPlabel.pop(imgidx)
- self.openNextImg()
- self.importDirImages(self.lastOpenDir, isDelete=True)
- def deleteImgDialog(self):
- yes, cancel = QMessageBox.Yes, QMessageBox.Cancel
- msg = u'The image will be deleted to the recycle bin'
- return QMessageBox.warning(self, u'Attention', msg, yes | cancel)
- def resetAll(self):
- self.settings.reset()
- self.close()
- proc = QProcess()
- proc.startDetached(os.path.abspath(__file__))
- def mayContinue(self): #
- if not self.dirty:
- return True
- else:
- discardChanges = self.discardChangesDialog()
- if discardChanges == QMessageBox.No:
- return True
- elif discardChanges == QMessageBox.Yes:
- self.saveFile()
- return True
- else:
- return False
- def discardChangesDialog(self):
- yes, no, cancel = QMessageBox.Yes, QMessageBox.No, QMessageBox.Cancel
- msg = u'You have unsaved changes, would you like to save them and proceed?\nClick "No" to undo all changes.'
- return QMessageBox.warning(self, u'Attention', msg, yes | no | cancel)
- def errorMessage(self, title, message):
- return QMessageBox.critical(self, title,
- '<p><b>%s</b></p>%s' % (title, message))
- def currentPath(self):
- return os.path.dirname(self.filePath) if self.filePath else '.'
- def chooseColor1(self):
- color = self.colorDialog.getColor(self.lineColor, u'Choose line color',
- default=DEFAULT_LINE_COLOR)
- if color:
- self.lineColor = color
- Shape.line_color = color
- self.canvas.setDrawingColor(color)
- self.canvas.update()
- self.setDirty()
- def deleteSelectedShape(self):
- self.remLabel(self.canvas.deleteSelected())
- self.setDirty()
- if self.noShapes():
- for action in self.actions.onShapesPresent:
- action.setEnabled(False)
- def chshapeLineColor(self):
- color = self.colorDialog.getColor(self.lineColor, u'Choose line color',
- default=DEFAULT_LINE_COLOR)
- if color:
- self.canvas.selectedShape.line_color = color
- self.canvas.update()
- self.setDirty()
- def chshapeFillColor(self):
- color = self.colorDialog.getColor(self.fillColor, u'Choose fill color',
- default=DEFAULT_FILL_COLOR)
- if color:
- self.canvas.selectedShape.fill_color = color
- self.canvas.update()
- self.setDirty()
- def copyShape(self):
- self.canvas.endMove(copy=True)
- self.addLabel(self.canvas.selectedShape)
- self.setDirty()
- def moveShape(self):
- self.canvas.endMove(copy=False)
- self.setDirty()
- def loadPredefinedClasses(self, predefClassesFile):
- if os.path.exists(predefClassesFile) is True:
- with codecs.open(predefClassesFile, 'r', 'utf8') as f:
- for line in f:
- line = line.strip()
- if self.labelHist is None:
- self.labelHist = [line]
- else:
- self.labelHist.append(line)
- def togglePaintLabelsOption(self):
- for shape in self.canvas.shapes:
- shape.paintLabel = self.displayLabelOption.isChecked()
- def toogleDrawSquare(self):
- self.canvas.setDrawingShapeToSquare(self.drawSquaresOption.isChecked())
- def additems(self, dirpath):
- for file in self.mImgList:
- pix = QPixmap(file)
- _, filename = os.path.split(file)
- filename, _ = os.path.splitext(filename)
- item = QListWidgetItem(QIcon(pix.scaled(100, 100, Qt.IgnoreAspectRatio, Qt.FastTransformation)),
- filename[:10])
- item.setToolTip(file)
- self.iconlist.addItem(item)
- def additems5(self, dirpath):
- for file in self.mImgList5:
- pix = QPixmap(file)
- _, filename = os.path.split(file)
- filename, _ = os.path.splitext(filename)
- pfilename = filename[:10]
- if len(pfilename) < 10:
- lentoken = 12 - len(pfilename)
- prelen = lentoken // 2
- bfilename = prelen * " " + pfilename + (lentoken - prelen) * " "
- # item = QListWidgetItem(QIcon(pix.scaled(100, 100, Qt.KeepAspectRatio, Qt.SmoothTransformation)),filename[:10])
- item = QListWidgetItem(QIcon(pix.scaled(100, 100, Qt.IgnoreAspectRatio, Qt.FastTransformation)),pfilename)
- # item.setForeground(QBrush(Qt.white))
- item.setToolTip(file)
- self.iconlist.addItem(item)
- owidth = 0
- for index in range(len(self.mImgList5)):
- item = self.iconlist.item(index)
- itemwidget = self.iconlist.visualItemRect(item)
- owidth += itemwidget.width()
- self.iconlist.setMinimumWidth(owidth + 50)
- def getImglabelidx(self, filePath):
- if platform.system()=='Windows':
- spliter = '\\'
- else:
- spliter = '/'
- filepathsplit = filePath.split(spliter)[-2:]
- return filepathsplit[0] + '/' + filepathsplit[1]
- def autoRecognition(self):
- assert self.mImgList is not None
- print('Using model from ', self.model)
- uncheckedList = [i for i in self.mImgList if i not in self.fileStatedict.keys()]
- self.autoDialog = AutoDialog(parent=self, ocr=self.ocr, mImgList=uncheckedList, lenbar=len(uncheckedList))
- self.autoDialog.popUp()
- self.currIndex=len(self.mImgList)
- self.loadFile(self.filePath) # ADD
- self.haveAutoReced = True
- self.AutoRecognition.setEnabled(False)
- self.actions.AutoRec.setEnabled(False)
- self.setDirty()
- self.saveCacheLabel()
- def reRecognition(self):
- img = cv2.imread(self.filePath)
- # org_box = [dic['points'] for dic in self.PPlabel[self.getImglabelidx(self.filePath)]]
- if self.canvas.shapes:
- self.result_dic = []
- rec_flag = 0
- for shape in self.canvas.shapes:
- box = [[int(p.x()), int(p.y())] for p in shape.points]
- assert len(box) == 4
- img_crop = get_rotate_crop_image(img, np.array(box, np.float32))
- if img_crop is None:
- msg = 'Can not recognise the detection box in ' + self.filePath + '. Please change manually'
- QMessageBox.information(self, "Information", msg)
- return
- result = self.ocr.ocr(img_crop, cls=True, det=False)
- if result[0][0] != '':
- result.insert(0, box)
- print('result in reRec is ', result)
- self.result_dic.append(result)
- if result[1][0] == shape.label:
- print('label no change')
- else:
- rec_flag += 1
- if len(self.result_dic) > 0 and rec_flag > 0:
- self.saveFile(mode='Auto')
- self.loadFile(self.filePath)
- self.setDirty()
- elif len(self.result_dic) == len(self.canvas.shapes) and rec_flag == 0:
- QMessageBox.information(self, "Information", "The recognition result remains unchanged!")
- else:
- print('Can not recgonise in ', self.filePath)
- else:
- QMessageBox.information(self, "Information", "Draw a box!")
- def singleRerecognition(self):
- img = cv2.imread(self.filePath)
- shape = self.canvas.selectedShape
- box = [[int(p.x()), int(p.y())] for p in shape.points]
- assert len(box) == 4
- img_crop = get_rotate_crop_image(img, np.array(box, np.float32))
- if img_crop is None:
- msg = 'Can not recognise the detection box in ' + self.filePath + '. Please change manually'
- QMessageBox.information(self, "Information", msg)
- return
- result = self.ocr.ocr(img_crop, cls=True, det=False)
- if result[0][0] != '':
- result.insert(0, box)
- print('result in reRec is ', result)
- if result[1][0] == shape.label:
- print('label no change')
- else:
- shape.label = result[1][0]
- self.singleLabel(shape)
- self.setDirty()
- print(box)
- def autolcm(self):
- vbox = QVBoxLayout()
- hbox = QHBoxLayout()
- self.panel = QLabel()
- self.panel.setText(self.stringBundle.getString('choseModelLg'))
- self.panel.setAlignment(Qt.AlignLeft)
- self.comboBox = QComboBox()
- self.comboBox.setObjectName("comboBox")
- self.comboBox.addItems(['Chinese & English', 'English', 'French', 'German', 'Korean', 'Japanese'])
- vbox.addWidget(self.panel)
- vbox.addWidget(self.comboBox)
- self.dialog = QDialog()
- self.dialog.resize(300, 100)
- self.okBtn = QPushButton(self.stringBundle.getString('ok'))
- self.cancelBtn = QPushButton(self.stringBundle.getString('cancel'))
- self.okBtn.clicked.connect(self.modelChoose)
- self.cancelBtn.clicked.connect(self.cancel)
- self.dialog.setWindowTitle(self.stringBundle.getString('choseModelLg'))
- hbox.addWidget(self.okBtn)
- hbox.addWidget(self.cancelBtn)
- vbox.addWidget(self.panel)
- vbox.addLayout(hbox)
- self.dialog.setLayout(vbox)
- self.dialog.setWindowModality(Qt.ApplicationModal)
- self.dialog.exec_()
- if self.filePath:
- self.AutoRecognition.setEnabled(True)
- self.actions.AutoRec.setEnabled(True)
- def modelChoose(self):
- print(self.comboBox.currentText())
- lg_idx = {'Chinese & English': 'ch', 'English': 'en', 'French': 'french', 'German': 'german',
- 'Korean': 'korean', 'Japanese': 'japan'}
- del self.ocr
- self.ocr = PaddleOCR(use_pdserving=False, use_angle_cls=True, det=True, cls=True, use_gpu=False,
- lang=lg_idx[self.comboBox.currentText()])
- self.dialog.close()
- def cancel(self):
- self.dialog.close()
- def loadFilestate(self, saveDir):
- self.fileStatepath = saveDir + '/fileState.txt'
- self.fileStatedict = {}
- if not os.path.exists(self.fileStatepath):
- f = open(self.fileStatepath, 'w', encoding='utf-8')
- else:
- with open(self.fileStatepath, 'r', encoding='utf-8') as f:
- states = f.readlines()
- for each in states:
- file, state = each.split('\t')
- self.fileStatedict[file] = 1
- self.actions.saveLabel.setEnabled(True)
- self.actions.saveRec.setEnabled(True)
- def saveFilestate(self):
- with open(self.fileStatepath, 'w', encoding='utf-8') as f:
- for key in self.fileStatedict:
- f.write(key + '\t')
- f.write(str(self.fileStatedict[key]) + '\n')
- def loadLabelFile(self, labelpath):
- labeldict = {}
- if not os.path.exists(labelpath):
- f = open(labelpath, 'w', encoding='utf-8')
- else:
- with open(labelpath, 'r', encoding='utf-8') as f:
- data = f.readlines()
- for each in data:
- file, label = each.split('\t')
- if label:
- label = label.replace('false', 'False')
- label = label.replace('true', 'True')
- labeldict[file] = eval(label)
- else:
- labeldict[file] = []
- return labeldict
- def savePPlabel(self,mode='Manual'):
- savedfile = [self.getImglabelidx(i) for i in self.fileStatedict.keys()]
- with open(self.PPlabelpath, 'w', encoding='utf-8') as f:
- for key in self.PPlabel:
- if key in savedfile and self.PPlabel[key] != []:
- f.write(key + '\t')
- f.write(json.dumps(self.PPlabel[key], ensure_ascii=False) + '\n')
- if mode=='Manual':
- msg = 'Images that have been checked are saved in '+ self.PPlabelpath
- QMessageBox.information(self, "Information", msg)
- def saveCacheLabel(self):
- with open(self.Cachelabelpath, 'w', encoding='utf-8') as f:
- for key in self.Cachelabel:
- f.write(key + '\t')
- f.write(json.dumps(self.Cachelabel[key], ensure_ascii=False) + '\n')
- def saveLabelFile(self):
- self.saveFilestate()
- self.savePPlabel()
- def saveRecResult(self):
- if None in [self.PPlabelpath, self.PPlabel, self.fileStatedict]:
- QMessageBox.information(self, "Information", "Save file first")
- return
- rec_gt_dir = os.path.dirname(self.PPlabelpath) + '/rec_gt.txt'
- crop_img_dir = os.path.dirname(self.PPlabelpath) + '/crop_img/'
- ques_img = []
- if not os.path.exists(crop_img_dir):
- os.mkdir(crop_img_dir)
- with open(rec_gt_dir, 'w', encoding='utf-8') as f:
- for key in self.fileStatedict:
- idx = self.getImglabelidx(key)
- try:
- img = cv2.imread(key)
- for i, label in enumerate(self.PPlabel[idx]):
- if label['difficult']: continue
- img_crop = get_rotate_crop_image(img, np.array(label['points'], np.float32))
- img_name = os.path.splitext(os.path.basename(idx))[0] + '_crop_'+str(i)+'.jpg'
- cv2.imwrite(crop_img_dir+img_name, img_crop)
- f.write('crop_img/'+ img_name + '\t')
- f.write(label['transcription'] + '\n')
- except Exception as e:
- ques_img.append(key)
- print("Can not read image ",e)
- if ques_img:
- QMessageBox.information(self, "Information", "The following images can not be saved, "
- "please check the image path and labels.\n" + "".join(str(i)+'\n' for i in ques_img))
- QMessageBox.information(self, "Information", "Cropped images have been saved in "+str(crop_img_dir))
- def speedChoose(self):
- if self.labelDialogOption.isChecked():
- self.canvas.newShape.disconnect()
- self.canvas.newShape.connect(partial(self.newShape, True))
- else:
- self.canvas.newShape.disconnect()
- self.canvas.newShape.connect(partial(self.newShape, False))
- def inverted(color):
- return QColor(*[255 - v for v in color.getRgb()])
- def read(filename, default=None):
- try:
- with open(filename, 'rb') as f:
- return f.read()
- except:
- return default
- def get_main_app(argv=[]):
- """
- Standard boilerplate Qt application code.
- Do everything but app.exec_() -- so that we can test the application in one thread
- """
- app = QApplication(argv)
- app.setApplicationName(__appname__)
- app.setWindowIcon(newIcon("app"))
- # Tzutalin 201705+: Accept extra agruments to change predefined class file
- argparser = argparse.ArgumentParser()
- argparser.add_argument("--lang", default='en', nargs="?")
- argparser.add_argument("--predefined_classes_file",
- default=os.path.join(os.path.dirname(__file__), "data", "predefined_classes.txt"),
- nargs="?")
- args = argparser.parse_args(argv[1:])
- # Usage : labelImg.py image predefClassFile saveDir
- win = MainWindow(lang=args.lang,
- defaultPrefdefClassFile=args.predefined_classes_file)
- win.show()
- return app, win
- def main():
- '''construct main app and run it'''
- app, _win = get_main_app(sys.argv)
- return app.exec_()
- if __name__ == '__main__':
-
- resource_file = './libs/resources.py'
- if not os.path.exists(resource_file):
- output = os.system('pyrcc5 -o libs/resources.py resources.qrc')
- assert output == 0, "operate the cmd have some problems ,please check whether there is a in the lib " \
- "directory resources.py "
- import libs.resources
- sys.exit(main())
|