2068 lines
82 KiB
Python
2068 lines
82 KiB
Python
# 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,
|
||
'', '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', getStr('creatPolygon'), 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)
|
||
|
||
undoLastPoint = action(getStr("undoLastPoint"), self.canvas.undoLastPoint,
|
||
'Ctrl+Z', "undo", getStr("undoLastPoint"), enabled=False)
|
||
|
||
undo = action(getStr("undo"), self.undoShapeEdit,
|
||
'Ctrl+Z', "undo", getStr("undo"), 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,
|
||
undo=undo, undoLastPoint=undoLastPoint,
|
||
fileMenuActions=(
|
||
opendir, saveLabel, resetAll, quit),
|
||
beginner=(), advanced=(),
|
||
editMenu=(createpoly, edit, copy, delete,singleRere,None, undo, undoLastPoint,
|
||
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)
|
||
|
||
self.autoSaveOption = QAction(getStr('autoSaveMode'), self)
|
||
self.autoSaveOption.setCheckable(True)
|
||
self.autoSaveOption.setChecked(settings.get(SETTING_PAINT_LABEL, False))
|
||
self.autoSaveOption.triggered.connect(self.autoSaveFunc)
|
||
|
||
addActions(self.menus.file,
|
||
(opendir, None, saveLabel, saveRec, self.autoSaveOption, 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)
|
||
self.actions.undoLastPoint.setEnabled(True)
|
||
|
||
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):
|
||
for shape in self.canvas.selectedShapes+[self.canvas.hShape]:
|
||
item = self.shapesToItemsbox[shape] # listitem
|
||
text = [(int(p.x()), int(p.y())) for p in shape.points]
|
||
item.setText(str(text))
|
||
self.actions.undo.setEnabled(True)
|
||
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())
|
||
|
||
|
||
def shapeSelectionChanged(self, selected_shapes):
|
||
self._noSelectionSlot = True
|
||
for shape in self.canvas.selectedShapes:
|
||
shape.selected = False
|
||
self.labelList.clearSelection()
|
||
self.canvas.selectedShapes = selected_shapes
|
||
for shape in self.canvas.selectedShapes:
|
||
shape.selected = True
|
||
self.shapesToItems[shape].setSelected(True)
|
||
self.shapesToItemsbox[shape].setSelected(True)
|
||
|
||
self.labelList.scrollToItem(self.currentItem()) # QAbstractItemView.EnsureVisible
|
||
self.BoxList.scrollToItem(self.currentBox())
|
||
|
||
self._noSelectionSlot = False
|
||
n_selected = len(selected_shapes)
|
||
self.actions.singleRere.setEnabled(n_selected)
|
||
self.actions.delete.setEnabled(n_selected)
|
||
self.actions.copy.setEnabled(n_selected)
|
||
self.actions.edit.setEnabled(n_selected == 1)
|
||
|
||
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 remLabels(self, shapes):
|
||
if shapes is None:
|
||
# print('rm empty label')
|
||
return
|
||
for shape in shapes:
|
||
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): # TODO：貌似没用
|
||
# 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):
|
||
for shape in self.canvas.copySelectedShape():
|
||
self.addLabel(shape)
|
||
# fix copy and delete
|
||
#self.shapeSelectionChanged(True)
|
||
|
||
|
||
def labelSelectionChanged(self):
|
||
if self._noSelectionSlot:
|
||
return
|
||
if self.canvas.editing():
|
||
selected_shapes = []
|
||
for item in self.labelList.selectedItems():
|
||
selected_shapes.append(self.itemsToShapes[item])
|
||
if selected_shapes:
|
||
self.canvas.selectShapes(selected_shapes)
|
||
else:
|
||
self.canvas.deSelectShape()
|
||
|
||
|
||
def boxSelectionChanged(self):
|
||
if self._noSelectionSlot:
|
||
#self.BoxList.scrollToItem(self.currentBox(), QAbstractItemView.PositionAtCenter)
|
||
return
|
||
if self.canvas.editing():
|
||
selected_shapes = []
|
||
for item in self.BoxList.selectedItems():
|
||
selected_shapes.append(self.itemsToShapesbox[item])
|
||
if selected_shapes:
|
||
self.canvas.selectShapes(selected_shapes)
|
||
else:
|
||
self.canvas.deSelectShape()
|
||
|
||
|
||
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)
|
||
self.actions.undoLastPoint.setEnabled(False)
|
||
self.actions.undo.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)
|
||
self.actions.saveLabel.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.remLabels(self.canvas.deleteSelected())
|
||
self.actions.undo.setEnabled(True)
|
||
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:
|
||
for shape in self.canvas.selectedShapes: shape.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:
|
||
for shape in self.canvas.selectedShapes: shape.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)
|
||
for shape in self.canvas.selectedShapes:
|
||
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 {} in [self.PPlabelpath, self.PPlabel, self.fileStatedict]:
|
||
QMessageBox.information(self, "Information", "Check the image 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 autoSaveFunc(self):
|
||
if self.autoSaveOption.isChecked():
|
||
self.autoSaveNum = 1 # Real auto_Save
|
||
try:
|
||
self.saveLabelFile()
|
||
except:
|
||
pass
|
||
print('The program will automatically save once after confirming an image')
|
||
else:
|
||
self.autoSaveNum = 5 # Used for backup
|
||
print('The program will automatically save once after confirming 5 images (default)')
|
||
|
||
def undoShapeEdit(self):
|
||
self.canvas.restoreShape()
|
||
self.labelList.clear()
|
||
self.BoxList.clear()
|
||
self.loadShapes(self.canvas.shapes)
|
||
self.actions.undo.setEnabled(self.canvas.isShapeRestorable)
|
||
|
||
def loadShapes(self, shapes, replace=True):
|
||
self._noSelectionSlot = True
|
||
for shape in shapes:
|
||
self.addLabel(shape)
|
||
self.labelList.clearSelection()
|
||
self._noSelectionSlot = False
|
||
self.canvas.loadShapes(shapes, replace=replace)
|
||
|
||
|
||
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())
|