diff --git a/PPOCRLabel/PPOCRLabel.py b/PPOCRLabel/PPOCRLabel.py index 517714104d1cb62f3b0c03c34843595d85502417..34c045e96aa10ba678447eefc1d007f9042804b8 100644 --- a/PPOCRLabel/PPOCRLabel.py +++ b/PPOCRLabel/PPOCRLabel.py @@ -10,7 +10,6 @@ # 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 @@ -24,13 +23,11 @@ import subprocess import sys from functools import partial -try: - from PyQt5 import QtCore, QtGui, QtWidgets - from PyQt5.QtGui import * - from PyQt5.QtCore import * - from PyQt5.QtWidgets import * -except ImportError: - print("Please install pyqt5...") +from PyQt5.QtCore import QSize, Qt, QPoint, QByteArray, QTimer, QFileInfo, QPointF, QProcess +from PyQt5.QtGui import QImage, QCursor, QPixmap, QImageReader +from PyQt5.QtWidgets import QMainWindow, QListWidget, QVBoxLayout, QToolButton, QHBoxLayout, QDockWidget, QWidget, \ + QSlider, QGraphicsOpacityEffect, QMessageBox, QListView, QScrollArea, QWidgetAction, QApplication, QLabel, \ + QFileDialog, QListWidgetItem, QComboBox, QDialog __dir__ = os.path.dirname(os.path.abspath(__file__)) @@ -42,6 +39,7 @@ sys.path.append("..") from paddleocr import PaddleOCR from libs.constants import * from libs.utils import * +from libs.labelColor import label_colormap from libs.settings import Settings from libs.shape import Shape, DEFAULT_LINE_COLOR, DEFAULT_FILL_COLOR, DEFAULT_LOCK_COLOR from libs.stringBundle import StringBundle @@ -53,9 +51,13 @@ from libs.colorDialog import ColorDialog from libs.ustr import ustr from libs.hashableQListWidgetItem import HashableQListWidgetItem from libs.editinlist import EditInList +from libs.unique_label_qlist_widget import UniqueLabelQListWidget +from libs.keyDialog import KeyDialog __appname__ = 'PPOCRLabel' +LABEL_COLORMAP = label_colormap() + class MainWindow(QMainWindow): FIT_WINDOW, FIT_WIDTH, MANUAL_ZOOM = list(range(3)) @@ -63,6 +65,7 @@ class MainWindow(QMainWindow): def __init__(self, lang="ch", gpu=False, + kie_mode=False, default_filename=None, default_predefined_class_file=None, default_save_dir=None): @@ -76,12 +79,19 @@ class MainWindow(QMainWindow): 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) + # KIE setting + self.kie_mode = kie_mode + self.key_previous_text = "" + self.existed_key_cls_set = set() + self.key_dialog_tip = getStr('keyDialogTip') + self.defaultSaveDir = default_save_dir self.ocr = PaddleOCR(use_pdserving=False, use_angle_cls=True, @@ -133,11 +143,13 @@ class MainWindow(QMainWindow): self.autoSaveNum = 5 # ================== File List ================== + + filelistLayout = QVBoxLayout() + filelistLayout.setContentsMargins(0, 0, 0, 0) + 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() @@ -158,10 +170,24 @@ class MainWindow(QMainWindow): self.fileDock.setWidget(fileListContainer) self.addDockWidget(Qt.LeftDockWidgetArea, self.fileDock) + # ================== Key List ================== + if self.kie_mode: + # self.keyList = QListWidget() + self.keyList = UniqueLabelQListWidget() + # self.keyList.itemSelectionChanged.connect(self.keyListSelectionChanged) + # self.keyList.itemDoubleClicked.connect(self.editBox) + # self.keyList.itemChanged.connect(self.keyListItemChanged) + self.keyListDockName = getStr('keyListTitle') + self.keyListDock = QDockWidget(self.keyListDockName, self) + self.keyListDock.setWidget(self.keyList) + self.keyListDock.setFeatures(QDockWidget.NoDockWidgetFeatures) + filelistLayout.addWidget(self.keyListDock) + # ================== Right Area ================== listLayout = QVBoxLayout() listLayout.setContentsMargins(0, 0, 0, 0) + # Buttons self.editButton = QToolButton() self.reRecogButton = QToolButton() self.reRecogButton.setIcon(newIcon('reRec', 30)) @@ -174,12 +200,12 @@ class MainWindow(QMainWindow): 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) + 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 @@ -341,7 +367,7 @@ class MainWindow(QMainWindow): resetAll = action(getStr('resetAll'), self.resetAll, None, 'resetall', getStr('resetAllDetail')) - color1 = action(getStr('boxLineColor'), self.chooseColor1, + color1 = action(getStr('boxLineColor'), self.chooseColor, 'Ctrl+L', 'color_line', getStr('boxLineColorDetail')) createMode = action(getStr('crtBox'), self.setCreateMode, @@ -402,11 +428,12 @@ class MainWindow(QMainWindow): self.MANUAL_ZOOM: lambda: 1, } + # ================== New Actions ================== + 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) @@ -437,6 +464,9 @@ class MainWindow(QMainWindow): undo = action(getStr("undo"), self.undoShapeEdit, 'Ctrl+Z', "undo", getStr("undo"), enabled=False) + change_cls = action(getStr("keyChange"), self.change_box_key, + 'Ctrl+B', "edit", getStr("keyChange"), enabled=False) + lock = action(getStr("lockBox"), self.lockSelectedShape, None, "lock", getStr("lockBoxDetail"), enabled=False) @@ -482,8 +512,7 @@ class MainWindow(QMainWindow): addActions(labelMenu, (edit, delete)) self.labelList.setContextMenuPolicy(Qt.CustomContextMenu) - self.labelList.customContextMenuRequested.connect( - self.popLabelListMenu) + self.labelList.customContextMenuRequested.connect(self.popLabelListMenu) # Draw squares/rectangles self.drawSquaresOption = QAction(getStr('drawSquares'), self) @@ -499,14 +528,15 @@ class MainWindow(QMainWindow): shapeLineColor=shapeLineColor, shapeFillColor=shapeFillColor, zoom=zoom, zoomIn=zoomIn, zoomOut=zoomOut, zoomOrg=zoomOrg, fitWindow=fitWindow, fitWidth=fitWidth, - zoomActions=zoomActions, saveLabel=saveLabel, + zoomActions=zoomActions, saveLabel=saveLabel, change_cls=change_cls, undo=undo, undoLastPoint=undoLastPoint, open_dataset_dir=open_dataset_dir, rotateLeft=rotateLeft, rotateRight=rotateRight, lock=lock, fileMenuActions=(opendir, open_dataset_dir, saveLabel, resetAll, quit), beginner=(), advanced=(), editMenu=(createpoly, edit, copy, delete, singleRere, None, undo, undoLastPoint, None, rotateLeft, rotateRight, None, color1, self.drawSquaresOption, lock), - beginnerContext=(create, edit, copy, delete, singleRere, rotateLeft, rotateRight, lock), + beginnerContext=( + create, edit, copy, delete, singleRere, rotateLeft, rotateRight, lock, change_cls), advancedContext=(createMode, editMode, edit, copy, delete, shapeLineColor, shapeFillColor), onLoadActive=(create, createMode, editMode), @@ -615,6 +645,8 @@ class MainWindow(QMainWindow): elif self.filePath: self.queueEvent(partial(self.loadFile, self.filePath or "")) + self.keyDialog = None + # Callbacks: self.zoomWidget.valueChanged.connect(self.paintCanvas) @@ -949,6 +981,12 @@ class MainWindow(QMainWindow): self.labelList.scrollToItem(self.currentItem()) # QAbstractItemView.EnsureVisible self.BoxList.scrollToItem(self.currentBox()) + if self.kie_mode: + if len(self.canvas.selectedShapes) == 1 and self.keyList.count() > 0: + selected_key_item_row = self.keyList.findItemsByLabel(self.canvas.selectedShapes[0].key_cls, + get_row=True) + self.keyList.setCurrentRow(selected_key_item_row) + self._noSelectionSlot = False n_selected = len(selected_shapes) self.actions.singleRere.setEnabled(n_selected) @@ -956,6 +994,7 @@ class MainWindow(QMainWindow): self.actions.copy.setEnabled(n_selected) self.actions.edit.setEnabled(n_selected == 1) self.actions.lock.setEnabled(n_selected) + self.actions.change_cls.setEnabled(n_selected) def addLabel(self, shape): shape.paintLabel = self.displayLabelOption.isChecked() @@ -1002,8 +1041,8 @@ class MainWindow(QMainWindow): def loadLabels(self, shapes): s = [] - for label, points, line_color, fill_color, difficult in shapes: - shape = Shape(label=label, line_color=line_color) + for label, points, line_color, key_cls, difficult in shapes: + shape = Shape(label=label, line_color=line_color, key_cls=key_cls) for x, y in points: # Ensure the labels are within the bounds of the image. If not, fix them. @@ -1017,16 +1056,7 @@ class MainWindow(QMainWindow): 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._update_shape_color(shape) self.addLabel(shape) self.updateComboBox() @@ -1066,14 +1096,16 @@ class MainWindow(QMainWindow): line_color=s.line_color.getRgb(), fill_color=s.fill_color.getRgb(), points=[(int(p.x()), int(p.y())) for p in s.points], # QPonitF - # add chris - difficult=s.difficult) # bool + difficult=s.difficult, + key_cls=s.key_cls) # bool - shapes = [] if mode == 'Auto' else \ - [format_shape(shape) for shape in self.canvas.shapes if shape.line_color != DEFAULT_LOCK_COLOR] + if mode == 'Auto': + shapes = [] + else: + shapes = [format_shape(shape) for shape in self.canvas.shapes if shape.line_color != DEFAULT_LOCK_COLOR] # Can add differrent annotation formats here for box in self.result_dic: - trans_dic = {"label": box[1][0], "points": box[0], 'difficult': False} + trans_dic = {"label": box[1][0], "points": box[0], "difficult": False, "key_cls": "None"} if trans_dic["label"] == "" and mode == 'Auto': continue shapes.append(trans_dic) @@ -1081,8 +1113,8 @@ class MainWindow(QMainWindow): try: trans_dic = [] for box in shapes: - trans_dic.append( - {"transcription": box['label'], "points": box['points'], 'difficult': box['difficult']}) + trans_dic.append({"transcription": box['label'], "points": box['points'], + "difficult": box['difficult'], "key_cls": box['key_cls']}) self.PPlabel[annotationFilePath] = trans_dic if mode == 'Auto': self.Cachelabel[annotationFilePath] = trans_dic @@ -1148,8 +1180,7 @@ class MainWindow(QMainWindow): position MUST be in global coordinates. """ if len(self.labelHist) > 0: - self.labelDialog = LabelDialog( - parent=self, listItem=self.labelHist) + self.labelDialog = LabelDialog(parent=self, listItem=self.labelHist) if value: text = self.labelDialog.popUp(text=self.prevLabelText) @@ -1159,8 +1190,22 @@ class MainWindow(QMainWindow): 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 + + shape = self.canvas.setLastLabel(text, None, None, None) # generate_color, generate_color + if self.kie_mode: + key_text, _ = self.keyDialog.popUp(self.key_previous_text) + if key_text is not None: + shape = self.canvas.setLastLabel(text, None, None, key_text) # generate_color, generate_color + self.key_previous_text = key_text + if not self.keyList.findItemsByLabel(key_text): + item = self.keyList.createItemFromLabel(key_text) + self.keyList.addItem(item) + rgb = self._get_rgb_by_label(key_text, self.kie_mode) + self.keyList.setItemLabel(item, key_text, rgb) + + self._update_shape_color(shape) + self.keyDialog.addLabelHistory(key_text) + self.addLabel(shape) if self.beginner(): # Switch to edit mode. self.canvas.setEditing(True) @@ -1175,6 +1220,25 @@ class MainWindow(QMainWindow): # self.canvas.undoLastLine() self.canvas.resetAllLines() + def _update_shape_color(self, shape): + r, g, b = self._get_rgb_by_label(shape.key_cls, self.kie_mode) + shape.line_color = QColor(r, g, b) + shape.vertex_fill_color = QColor(r, g, b) + shape.hvertex_fill_color = QColor(255, 255, 255) + shape.fill_color = QColor(r, g, b, 128) + shape.select_line_color = QColor(255, 255, 255) + shape.select_fill_color = QColor(r, g, b, 155) + + def _get_rgb_by_label(self, label, kie_mode): + shift_auto_shape_color = 2 # use for random color + if kie_mode and label != "None": + item = self.keyList.findItemsByLabel(label)[0] + label_id = self.keyList.indexFromItem(item).row() + 1 + label_id += shift_auto_shape_color + return LABEL_COLORMAP[label_id % len(LABEL_COLORMAP)] + else: + return (0, 255, 0) + def scrollRequest(self, delta, orientation): units = - delta / (8 * 15) bar = self.scrollBars[orientation] @@ -1344,7 +1408,7 @@ class MainWindow(QMainWindow): select_indexes = self.fileListWidget.selectedIndexes() if len(select_indexes) > 0: self.fileDock.setWindowTitle(self.fileListName + f" ({select_indexes[0].row() + 1}" - f"/{self.fileListWidget.count()})") + f"/{self.fileListWidget.count()})") # update show counting self.BoxListDock.setWindowTitle(self.BoxListDockName + f" ({self.BoxList.count()})") self.labelListDock.setWindowTitle(self.labelListDockName + f" ({self.labelList.count()})") @@ -1362,13 +1426,13 @@ class MainWindow(QMainWindow): for box in self.canvas.lockedShapes: if self.canvas.isInTheSameImage: shapes.append((box['transcription'], [[s[0] * width, s[1] * height] for s in box['ratio']], - DEFAULT_LOCK_COLOR, None, box['difficult'])) + DEFAULT_LOCK_COLOR, box['key_cls'], box['difficult'])) else: shapes.append(('锁定框:待检测', [[s[0] * width, s[1] * height] for s in box['ratio']], - DEFAULT_LOCK_COLOR, None, box['difficult'])) + DEFAULT_LOCK_COLOR, box['key_cls'], box['difficult'])) if imgidx in self.PPlabel.keys(): for box in self.PPlabel[imgidx]: - shapes.append((box['transcription'], box['points'], None, None, box['difficult'])) + shapes.append((box['transcription'], box['points'], None, box['key_cls'], box['difficult'])) self.loadLabels(shapes) self.canvas.verified = False @@ -1504,6 +1568,39 @@ class MainWindow(QMainWindow): self.actions.open_dataset_dir.setEnabled(False) defaultOpenDirPath = os.path.dirname(self.filePath) if self.filePath else '.' + def init_key_list(self, label_dict): + if not self.kie_mode: + return + # load key_cls + for image, info in label_dict.items(): + for box in info: + if "key_cls" not in box: + continue + self.existed_key_cls_set.add(box["key_cls"]) + if len(self.existed_key_cls_set) > 0: + for key_text in self.existed_key_cls_set: + if not self.keyList.findItemsByLabel(key_text): + item = self.keyList.createItemFromLabel(key_text) + self.keyList.addItem(item) + rgb = self._get_rgb_by_label(key_text, self.kie_mode) + self.keyList.setItemLabel(item, key_text, rgb) + + if self.keyDialog is None: + # key list dialog + self.keyDialog = KeyDialog( + text=self.key_dialog_tip, + parent=self, + labels=self.existed_key_cls_set, + sort_labels=True, + show_text_field=True, + completion="startswith", + fit_to_content={'column': True, 'row': False}, + flags=None + ) + else: + self.keyDialog.labelList.addItems(self.existed_key_cls_set) + + def importDirImages(self, dirpath, isDelete=False): if not self.mayContinue() or not dirpath: return @@ -1518,6 +1615,9 @@ class MainWindow(QMainWindow): self.Cachelabel = self.loadLabelFile(self.Cachelabelpath) if self.Cachelabel: self.PPlabel = dict(self.Cachelabel, **self.PPlabel) + + self.init_key_list(self.PPlabel) + self.lastOpenDir = dirpath self.dirname = dirpath @@ -1737,7 +1837,7 @@ class MainWindow(QMainWindow): def currentPath(self): return os.path.dirname(self.filePath) if self.filePath else '.' - def chooseColor1(self): + def chooseColor(self): color = self.colorDialog.getColor(self.lineColor, u'Choose line color', default=DEFAULT_LINE_COLOR) if color: @@ -1854,6 +1954,8 @@ class MainWindow(QMainWindow): self.setDirty() self.saveCacheLabel() + self.init_key_list(self.Cachelabel) + def reRecognition(self): img = cv2.imread(self.filePath) # org_box = [dic['points'] for dic in self.PPlabel[self.getImglabelidx(self.filePath)]] @@ -2059,7 +2161,8 @@ class MainWindow(QMainWindow): try: img = cv2.imread(key) for i, label in enumerate(self.PPlabel[idx]): - if label['difficult']: continue + 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) @@ -2096,6 +2199,15 @@ class MainWindow(QMainWindow): self.autoSaveNum = 5 # Used for backup print('The program will automatically save once after confirming 5 images (default)') + def change_box_key(self): + key_text, _ = self.keyDialog.popUp(self.key_previous_text) + if key_text is None: + return + self.key_previous_text = key_text + for shape in self.canvas.selectedShapes: + shape.key_cls = key_text + self._update_shape_color(shape) + def undoShapeEdit(self): self.canvas.restoreShape() self.labelList.clear() @@ -2126,8 +2238,9 @@ class MainWindow(QMainWindow): line_color=s.line_color.getRgb(), fill_color=s.fill_color.getRgb(), ratio=[[int(p.x()) / width, int(p.y()) / height] for p in s.points], # QPonitF - # add chris - difficult=s.difficult) # bool + difficult=s.difficult, # bool + key_cls=s.key_cls, # bool + ) # lock if len(self.canvas.lockedShapes) == 0: @@ -2137,7 +2250,9 @@ class MainWindow(QMainWindow): shapes = [format_shape(shape) for shape in self.canvas.selectedShapes] trans_dic = [] for box in shapes: - trans_dic.append({"transcription": box['label'], "ratio": box['ratio'], 'difficult': box['difficult']}) + trans_dic.append({"transcription": box['label'], "ratio": box['ratio'], + "difficult": box['difficult'], + "key_cls": "None" if "key_cls" not in box else box["key_cls"]}) self.canvas.lockedShapes = trans_dic self.actions.save.setEnabled(True) @@ -2179,6 +2294,7 @@ def get_main_app(argv=[]): arg_parser = argparse.ArgumentParser() arg_parser.add_argument("--lang", type=str, default='en', nargs="?") arg_parser.add_argument("--gpu", type=str2bool, default=True, nargs="?") + arg_parser.add_argument("--kie", type=str2bool, default=False, nargs="?") arg_parser.add_argument("--predefined_classes_file", default=os.path.join(os.path.dirname(__file__), "data", "predefined_classes.txt"), nargs="?") @@ -2186,6 +2302,7 @@ def get_main_app(argv=[]): win = MainWindow(lang=args.lang, gpu=args.gpu, + kie_mode=args.kie, default_predefined_class_file=args.predefined_classes_file) win.show() return app, win diff --git a/PPOCRLabel/README.md b/PPOCRLabel/README.md index e40d82916ed30f39641e88e73c6563cc4cd0183b..4d25e670ae6d07d569a247bc5f9c35c939b23f8e 100644 --- a/PPOCRLabel/README.md +++ b/PPOCRLabel/README.md @@ -8,6 +8,8 @@ PPOCRLabel is a semi-automatic graphic annotation tool suitable for OCR field, w ### Recent Update +- 2022.02:(by [PeterH0323](https://github.com/peterh0323) ) + - Added KIE mode, for [detection + identification + keyword extraction] labeling. - 2022.01:(by [PeterH0323](https://github.com/peterh0323) ) - Improve user experience: prompt for the number of files and labels, optimize interaction, and fix bugs such as only use CPU when inference - 2021.11.17: @@ -72,7 +74,8 @@ PPOCRLabel ```bash pip3 install PPOCRLabel pip3 install opencv-contrib-python-headless==4.2.0.32 -PPOCRLabel # run +PPOCRLabel # [Normal mode] for [detection + recognition] labeling +PPOCRLabel --kie True # [KIE mode] for [detection + recognition + keyword extraction] labeling ``` #### 1.2.2 Build and Install the Whl Package Locally @@ -87,7 +90,8 @@ pip3 install dist/PPOCRLabel-1.0.2-py2.py3-none-any.whl ```bash cd ./PPOCRLabel # Switch to the PPOCRLabel directory -python PPOCRLabel.py +python PPOCRLabel.py # [Normal mode] for [detection + recognition] labeling +python PPOCRLabel.py --kie True # [KIE mode] for [detection + recognition + keyword extraction] labeling ``` diff --git a/PPOCRLabel/README_ch.md b/PPOCRLabel/README_ch.md index 815de99729e00fef6be0247571e27aebc82e53ab..3f8dc4f0c6b7cc88a71409d123f598e74b3f2cad 100644 --- a/PPOCRLabel/README_ch.md +++ b/PPOCRLabel/README_ch.md @@ -8,6 +8,8 @@ PPOCRLabel是一款适用于OCR领域的半自动化图形标注工具,内置P #### 近期更新 +- 2022.02:(by [PeterH0323](https://github.com/peterh0323) ) + - 新增:KIE 功能,用于打【检测+识别+关键字提取】的标签 - 2022.01:(by [PeterH0323](https://github.com/peterh0323) ) - 提升用户体验:新增文件与标记数目提示、优化交互、修复gpu使用等问题 - 2021.11.17: @@ -70,7 +72,8 @@ PPOCRLabel --lang ch ```bash pip3 install PPOCRLabel pip3 install opencv-contrib-python-headless==4.2.0.32 # 如果下载过慢请添加"-i https://mirror.baidu.com/pypi/simple" -PPOCRLabel --lang ch # 启动 +PPOCRLabel --lang ch # 启动【普通模式】,用于打【检测+识别】场景的标签 +PPOCRLabel --lang ch --kie True # 启动 【KIE 模式】,用于打【检测+识别+关键字提取】场景的标签 ``` > 如果上述安装出现问题,可以参考3.6节 错误提示 @@ -89,7 +92,8 @@ pip3 install dist/PPOCRLabel-1.0.2-py2.py3-none-any.whl -i https://mirror.baidu. ```bash cd ./PPOCRLabel # 切换到PPOCRLabel目录 -python PPOCRLabel.py --lang ch +python PPOCRLabel.py --lang ch # 启动【普通模式】,用于打【检测+识别】场景的标签 +python PPOCRLabel.py --lang ch --kie True # 启动 【KIE 模式】,用于打【检测+识别+关键字提取】场景的标签 ``` diff --git a/PPOCRLabel/libs/canvas.py b/PPOCRLabel/libs/canvas.py index 8d257e6bd7e7a61d7c28e9787042c3eb9d42609f..095fe5ab06553dcb05c8bcc061f950ded606ebb3 100644 --- a/PPOCRLabel/libs/canvas.py +++ b/PPOCRLabel/libs/canvas.py @@ -783,7 +783,7 @@ class Canvas(QWidget): points = [p1+p2 for p1, p2 in zip(self.selectedShape.points, [step]*4)] return True in map(self.outOfPixmap, points) - def setLastLabel(self, text, line_color = None, fill_color = None): + def setLastLabel(self, text, line_color=None, fill_color=None, key_cls=None): assert text self.shapes[-1].label = text if line_color: @@ -791,6 +791,10 @@ class Canvas(QWidget): if fill_color: self.shapes[-1].fill_color = fill_color + + if key_cls: + self.shapes[-1].key_cls = key_cls + self.storeShapes() return self.shapes[-1] diff --git a/PPOCRLabel/libs/keyDialog.py b/PPOCRLabel/libs/keyDialog.py new file mode 100644 index 0000000000000000000000000000000000000000..1ec8d97147cd2eb1e3c8482a9a6c5092edcd1b9c --- /dev/null +++ b/PPOCRLabel/libs/keyDialog.py @@ -0,0 +1,216 @@ +import re + +from PyQt5 import QtCore +from PyQt5 import QtGui +from PyQt5 import QtWidgets +from PyQt5.Qt import QT_VERSION_STR +from libs.utils import newIcon, labelValidator + +QT5 = QT_VERSION_STR[0] == '5' + + +# TODO(unknown): +# - Calculate optimal position so as not to go out of screen area. + + +class KeyQLineEdit(QtWidgets.QLineEdit): + def setListWidget(self, list_widget): + self.list_widget = list_widget + + def keyPressEvent(self, e): + if e.key() in [QtCore.Qt.Key_Up, QtCore.Qt.Key_Down]: + self.list_widget.keyPressEvent(e) + else: + super(KeyQLineEdit, self).keyPressEvent(e) + + +class KeyDialog(QtWidgets.QDialog): + def __init__( + self, + text="Enter object label", + parent=None, + labels=None, + sort_labels=True, + show_text_field=True, + completion="startswith", + fit_to_content=None, + flags=None, + ): + if fit_to_content is None: + fit_to_content = {"row": False, "column": True} + self._fit_to_content = fit_to_content + + super(KeyDialog, self).__init__(parent) + self.edit = KeyQLineEdit() + self.edit.setPlaceholderText(text) + self.edit.setValidator(labelValidator()) + self.edit.editingFinished.connect(self.postProcess) + if flags: + self.edit.textChanged.connect(self.updateFlags) + + layout = QtWidgets.QVBoxLayout() + if show_text_field: + layout_edit = QtWidgets.QHBoxLayout() + layout_edit.addWidget(self.edit, 6) + layout.addLayout(layout_edit) + # buttons + self.buttonBox = bb = QtWidgets.QDialogButtonBox( + QtWidgets.QDialogButtonBox.Ok | QtWidgets.QDialogButtonBox.Cancel, + QtCore.Qt.Horizontal, + self, + ) + bb.button(bb.Ok).setIcon(newIcon("done")) + bb.button(bb.Cancel).setIcon(newIcon("undo")) + bb.accepted.connect(self.validate) + bb.rejected.connect(self.reject) + layout.addWidget(bb) + # label_list + self.labelList = QtWidgets.QListWidget() + if self._fit_to_content["row"]: + self.labelList.setHorizontalScrollBarPolicy( + QtCore.Qt.ScrollBarAlwaysOff + ) + if self._fit_to_content["column"]: + self.labelList.setVerticalScrollBarPolicy( + QtCore.Qt.ScrollBarAlwaysOff + ) + self._sort_labels = sort_labels + if labels: + self.labelList.addItems(labels) + if self._sort_labels: + self.labelList.sortItems() + else: + self.labelList.setDragDropMode( + QtWidgets.QAbstractItemView.InternalMove + ) + self.labelList.currentItemChanged.connect(self.labelSelected) + self.labelList.itemDoubleClicked.connect(self.labelDoubleClicked) + self.edit.setListWidget(self.labelList) + layout.addWidget(self.labelList) + # label_flags + if flags is None: + flags = {} + self._flags = flags + self.flagsLayout = QtWidgets.QVBoxLayout() + self.resetFlags() + layout.addItem(self.flagsLayout) + self.edit.textChanged.connect(self.updateFlags) + self.setLayout(layout) + # completion + completer = QtWidgets.QCompleter() + if not QT5 and completion != "startswith": + completion = "startswith" + if completion == "startswith": + completer.setCompletionMode(QtWidgets.QCompleter.InlineCompletion) + # Default settings. + # completer.setFilterMode(QtCore.Qt.MatchStartsWith) + elif completion == "contains": + completer.setCompletionMode(QtWidgets.QCompleter.PopupCompletion) + completer.setFilterMode(QtCore.Qt.MatchContains) + else: + raise ValueError("Unsupported completion: {}".format(completion)) + completer.setModel(self.labelList.model()) + self.edit.setCompleter(completer) + + def addLabelHistory(self, label): + if self.labelList.findItems(label, QtCore.Qt.MatchExactly): + return + self.labelList.addItem(label) + if self._sort_labels: + self.labelList.sortItems() + + def labelSelected(self, item): + self.edit.setText(item.text()) + + def validate(self): + text = self.edit.text() + if hasattr(text, "strip"): + text = text.strip() + else: + text = text.trimmed() + if text: + self.accept() + + def labelDoubleClicked(self, item): + self.validate() + + def postProcess(self): + text = self.edit.text() + if hasattr(text, "strip"): + text = text.strip() + else: + text = text.trimmed() + self.edit.setText(text) + + def updateFlags(self, label_new): + # keep state of shared flags + flags_old = self.getFlags() + + flags_new = {} + for pattern, keys in self._flags.items(): + if re.match(pattern, label_new): + for key in keys: + flags_new[key] = flags_old.get(key, False) + self.setFlags(flags_new) + + def deleteFlags(self): + for i in reversed(range(self.flagsLayout.count())): + item = self.flagsLayout.itemAt(i).widget() + self.flagsLayout.removeWidget(item) + item.setParent(None) + + def resetFlags(self, label=""): + flags = {} + for pattern, keys in self._flags.items(): + if re.match(pattern, label): + for key in keys: + flags[key] = False + self.setFlags(flags) + + def setFlags(self, flags): + self.deleteFlags() + for key in flags: + item = QtWidgets.QCheckBox(key, self) + item.setChecked(flags[key]) + self.flagsLayout.addWidget(item) + item.show() + + def getFlags(self): + flags = {} + for i in range(self.flagsLayout.count()): + item = self.flagsLayout.itemAt(i).widget() + flags[item.text()] = item.isChecked() + return flags + + def popUp(self, text=None, move=True, flags=None): + if self._fit_to_content["row"]: + self.labelList.setMinimumHeight( + self.labelList.sizeHintForRow(0) * self.labelList.count() + 2 + ) + if self._fit_to_content["column"]: + self.labelList.setMinimumWidth( + self.labelList.sizeHintForColumn(0) + 2 + ) + # if text is None, the previous label in self.edit is kept + if text is None: + text = self.edit.text() + if flags: + self.setFlags(flags) + else: + self.resetFlags(text) + self.edit.setText(text) + self.edit.setSelection(0, len(text)) + + items = self.labelList.findItems(text, QtCore.Qt.MatchFixedString) + if items: + if len(items) != 1: + self.labelList.setCurrentItem(items[0]) + row = self.labelList.row(items[0]) + self.edit.completer().setCurrentRow(row) + self.edit.setFocus(QtCore.Qt.PopupFocusReason) + if move: + self.move(QtGui.QCursor.pos()) + if self.exec_(): + return self.edit.text(), self.getFlags() + else: + return None, None diff --git a/PPOCRLabel/libs/labelColor.py b/PPOCRLabel/libs/labelColor.py new file mode 100644 index 0000000000000000000000000000000000000000..c6f933981f3ca13981910a88fca76f884d727a14 --- /dev/null +++ b/PPOCRLabel/libs/labelColor.py @@ -0,0 +1,88 @@ +import PIL.Image +import numpy as np + + +def rgb2hsv(rgb): + # type: (np.ndarray) -> np.ndarray + """Convert rgb to hsv. + + Parameters + ---------- + rgb: numpy.ndarray, (H, W, 3), np.uint8 + Input rgb image. + + Returns + ------- + hsv: numpy.ndarray, (H, W, 3), np.uint8 + Output hsv image. + + """ + hsv = PIL.Image.fromarray(rgb, mode="RGB") + hsv = hsv.convert("HSV") + hsv = np.array(hsv) + return hsv + + +def hsv2rgb(hsv): + # type: (np.ndarray) -> np.ndarray + """Convert hsv to rgb. + + Parameters + ---------- + hsv: numpy.ndarray, (H, W, 3), np.uint8 + Input hsv image. + + Returns + ------- + rgb: numpy.ndarray, (H, W, 3), np.uint8 + Output rgb image. + + """ + rgb = PIL.Image.fromarray(hsv, mode="HSV") + rgb = rgb.convert("RGB") + rgb = np.array(rgb) + return rgb + + +def label_colormap(n_label=256, value=None): + """Label colormap. + + Parameters + ---------- + n_label: int + Number of labels (default: 256). + value: float or int + Value scale or value of label color in HSV space. + + Returns + ------- + cmap: numpy.ndarray, (N, 3), numpy.uint8 + Label id to colormap. + + """ + + def bitget(byteval, idx): + return (byteval & (1 << idx)) != 0 + + cmap = np.zeros((n_label, 3), dtype=np.uint8) + for i in range(0, n_label): + id = i + r, g, b = 0, 0, 0 + for j in range(0, 8): + r = np.bitwise_or(r, (bitget(id, 0) << 7 - j)) + g = np.bitwise_or(g, (bitget(id, 1) << 7 - j)) + b = np.bitwise_or(b, (bitget(id, 2) << 7 - j)) + id = id >> 3 + cmap[i, 0] = r + cmap[i, 1] = g + cmap[i, 2] = b + + if value is not None: + hsv = rgb2hsv(cmap.reshape(1, -1, 3)) + if isinstance(value, float): + hsv[:, 1:, 2] = hsv[:, 1:, 2].astype(float) * value + else: + assert isinstance(value, int) + hsv[:, 1:, 2] = value + cmap = hsv2rgb(hsv).reshape(-1, 3) + return cmap diff --git a/PPOCRLabel/libs/shape.py b/PPOCRLabel/libs/shape.py index 528b1102b010ceef8fa1057309e652010a91376d..fc8ab5ec4d7ff2836034d9c7e01acaf49dfe7aa0 100644 --- a/PPOCRLabel/libs/shape.py +++ b/PPOCRLabel/libs/shape.py @@ -46,12 +46,13 @@ class Shape(object): point_size = 8 scale = 1.0 - def __init__(self, label=None, line_color=None, difficult=False, paintLabel=False): + def __init__(self, label=None, line_color=None, difficult=False, key_cls="None", paintLabel=False): self.label = label self.points = [] self.fill = False self.selected = False self.difficult = difficult + self.key_cls = key_cls self.paintLabel = paintLabel self.locked = False self.direction = 0 @@ -224,6 +225,7 @@ class Shape(object): if self.fill_color != Shape.fill_color: shape.fill_color = self.fill_color shape.difficult = self.difficult + shape.key_cls = self.key_cls return shape def __len__(self): diff --git a/PPOCRLabel/libs/unique_label_qlist_widget.py b/PPOCRLabel/libs/unique_label_qlist_widget.py new file mode 100644 index 0000000000000000000000000000000000000000..f1eff7a172d3fecf9c18579ccead5f62ba65ecd5 --- /dev/null +++ b/PPOCRLabel/libs/unique_label_qlist_widget.py @@ -0,0 +1,45 @@ +# -*- encoding: utf-8 -*- + +from PyQt5.QtCore import Qt +from PyQt5 import QtWidgets + + +class EscapableQListWidget(QtWidgets.QListWidget): + def keyPressEvent(self, event): + super(EscapableQListWidget, self).keyPressEvent(event) + if event.key() == Qt.Key_Escape: + self.clearSelection() + + +class UniqueLabelQListWidget(EscapableQListWidget): + def mousePressEvent(self, event): + super(UniqueLabelQListWidget, self).mousePressEvent(event) + if not self.indexAt(event.pos()).isValid(): + self.clearSelection() + + def findItemsByLabel(self, label, get_row=False): + items = [] + for row in range(self.count()): + item = self.item(row) + if item.data(Qt.UserRole) == label: + items.append(item) + if get_row: + return row + return items + + def createItemFromLabel(self, label): + item = QtWidgets.QListWidgetItem() + item.setData(Qt.UserRole, label) + return item + + def setItemLabel(self, item, label, color=None): + qlabel = QtWidgets.QLabel() + if color is None: + qlabel.setText(f"{label}") + else: + qlabel.setText(' {} '.format(*color, label)) + qlabel.setAlignment(Qt.AlignBottom) + + item.setSizeHint(qlabel.sizeHint()) + + self.setItemWidget(item, qlabel) diff --git a/PPOCRLabel/libs/utils.py b/PPOCRLabel/libs/utils.py index 9fab41d3ffee33b8f86f9576507eb13b18806496..2510520caa8048d7787d7c8f65df2885d76026f7 100644 --- a/PPOCRLabel/libs/utils.py +++ b/PPOCRLabel/libs/utils.py @@ -10,30 +10,26 @@ # 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. -from math import sqrt -from libs.ustr import ustr import hashlib +import os import re import sys +from math import sqrt + import cv2 import numpy as np -import os +from PyQt5.QtCore import QRegExp, QT_VERSION_STR +from PyQt5.QtGui import QIcon, QRegExpValidator, QColor +from PyQt5.QtWidgets import QPushButton, QAction, QMenu +from libs.ustr import ustr -__dir__ = os.path.dirname(os.path.abspath(__file__)) # 获取本程序文件路径 +__dir__ = os.path.dirname(os.path.abspath(__file__)) # 获取本程序文件路径 __iconpath__ = os.path.abspath(os.path.join(__dir__, '../resources/icons')) -try: - from PyQt5.QtGui import * - from PyQt5.QtCore import * - from PyQt5.QtWidgets import * -except ImportError: - from PyQt4.QtGui import * - from PyQt4.QtCore import * - def newIcon(icon, iconSize=None): if iconSize is not None: - return QIcon(QIcon(__iconpath__ + "/" + icon + ".png").pixmap(iconSize,iconSize)) + return QIcon(QIcon(__iconpath__ + "/" + icon + ".png").pixmap(iconSize, iconSize)) else: return QIcon(__iconpath__ + "/" + icon + ".png") @@ -105,24 +101,25 @@ def generateColorByText(text): s = ustr(text) hashCode = int(hashlib.sha256(s.encode('utf-8')).hexdigest(), 16) r = int((hashCode / 255) % 255) - g = int((hashCode / 65025) % 255) - b = int((hashCode / 16581375) % 255) + g = int((hashCode / 65025) % 255) + b = int((hashCode / 16581375) % 255) return QColor(r, g, b, 100) + def have_qstring(): '''p3/qt5 get rid of QString wrapper as py3 has native unicode str type''' return not (sys.version_info.major >= 3 or QT_VERSION_STR.startswith('5.')) -def util_qt_strlistclass(): - return QStringList if have_qstring() else list -def natural_sort(list, key=lambda s:s): +def natural_sort(list, key=lambda s: s): """ Sort the list into natural alphanumeric order. """ + def get_alphanum_key_func(key): convert = lambda text: int(text) if text.isdigit() else text return lambda s: [convert(c) for c in re.split('([0-9]+)', key(s))] + sort_key = get_alphanum_key_func(key) list.sort(key=sort_key) @@ -133,8 +130,8 @@ def get_rotate_crop_image(img, points): d = 0.0 for index in range(-1, 3): d += -0.5 * (points[index + 1][1] + points[index][1]) * ( - points[index + 1][0] - points[index][0]) - if d < 0: # counterclockwise + points[index + 1][0] - points[index][0]) + if d < 0: # counterclockwise tmp = np.array(points) points[1], points[3] = tmp[3], tmp[1] @@ -163,10 +160,11 @@ def get_rotate_crop_image(img, points): except Exception as e: print(e) + def stepsInfo(lang='en'): if lang == 'ch': msg = "1. 安装与运行:使用上述命令安装与运行程序。\n" \ - "2. 打开文件夹:在菜单栏点击 “文件” - 打开目录 选择待标记图片的文件夹.\n"\ + "2. 打开文件夹:在菜单栏点击 “文件” - 打开目录 选择待标记图片的文件夹.\n" \ "3. 自动标注:点击 ”自动标注“,使用PPOCR超轻量模型对图片文件名前图片状态为 “X” 的图片进行自动标注。\n" \ "4. 手动标注:点击 “矩形标注”(推荐直接在英文模式下点击键盘中的 “W”),用户可对当前图片中模型未检出的部分进行手动" \ "绘制标记框。点击键盘P,则使用四点标注模式(或点击“编辑” - “四点标注”),用户依次点击4个点后,双击左键表示标注完成。\n" \ @@ -181,25 +179,26 @@ def stepsInfo(lang='en'): else: msg = "1. Build and launch using the instructions above.\n" \ - "2. Click 'Open Dir' in Menu/File to select the folder of the picture.\n"\ - "3. Click 'Auto recognition', use PPOCR model to automatically annotate images which marked with 'X' before the file name."\ - "4. Create Box:\n"\ - "4.1 Click 'Create RectBox' or press 'W' in English keyboard mode to draw a new rectangle detection box. Click and release left mouse to select a region to annotate the text area.\n"\ - "4.2 Press 'P' to enter four-point labeling mode which enables you to create any four-point shape by clicking four points with the left mouse button in succession and DOUBLE CLICK the left mouse as the signal of labeling completion.\n"\ - "5. After the marking frame is drawn, the user clicks 'OK', and the detection frame will be pre-assigned a TEMPORARY label.\n"\ - "6. Click re-Recognition, model will rewrite ALL recognition results in ALL detection box.\n"\ - "7. Double click the result in 'recognition result' list to manually change inaccurate recognition results.\n"\ - "8. Click 'Save', the image status will switch to '√',then the program automatically jump to the next.\n"\ - "9. Click 'Delete Image' and the image will be deleted to the recycle bin.\n"\ - "10. Labeling result: After closing the application or switching the file path, the manually saved label will be stored in *Label.txt* under the opened picture folder.\n"\ - " Click PaddleOCR-Save Recognition Results in the menu bar, the recognition training data of such pictures will be saved in the *crop_img* folder, and the recognition label will be saved in *rec_gt.txt*.\n" + "2. Click 'Open Dir' in Menu/File to select the folder of the picture.\n" \ + "3. Click 'Auto recognition', use PPOCR model to automatically annotate images which marked with 'X' before the file name." \ + "4. Create Box:\n" \ + "4.1 Click 'Create RectBox' or press 'W' in English keyboard mode to draw a new rectangle detection box. Click and release left mouse to select a region to annotate the text area.\n" \ + "4.2 Press 'P' to enter four-point labeling mode which enables you to create any four-point shape by clicking four points with the left mouse button in succession and DOUBLE CLICK the left mouse as the signal of labeling completion.\n" \ + "5. After the marking frame is drawn, the user clicks 'OK', and the detection frame will be pre-assigned a TEMPORARY label.\n" \ + "6. Click re-Recognition, model will rewrite ALL recognition results in ALL detection box.\n" \ + "7. Double click the result in 'recognition result' list to manually change inaccurate recognition results.\n" \ + "8. Click 'Save', the image status will switch to '√',then the program automatically jump to the next.\n" \ + "9. Click 'Delete Image' and the image will be deleted to the recycle bin.\n" \ + "10. Labeling result: After closing the application or switching the file path, the manually saved label will be stored in *Label.txt* under the opened picture folder.\n" \ + " Click PaddleOCR-Save Recognition Results in the menu bar, the recognition training data of such pictures will be saved in the *crop_img* folder, and the recognition label will be saved in *rec_gt.txt*.\n" return msg + def keysInfo(lang='en'): if lang == 'ch': msg = "快捷键\t\t\t说明\n" \ - "———————————————————————\n"\ + "———————————————————————\n" \ "Ctrl + shift + R\t\t对当前图片的所有标记重新识别\n" \ "W\t\t\t新建矩形框\n" \ "Q\t\t\t新建四点框\n" \ @@ -223,17 +222,17 @@ def keysInfo(lang='en'): "———————————————————————\n" \ "Ctrl + shift + R\t\tRe-recognize all the labels\n" \ "\t\t\tof the current image\n" \ - "\n"\ + "\n" \ "W\t\t\tCreate a rect box\n" \ "Q\t\t\tCreate a four-points box\n" \ "Ctrl + E\t\tEdit label of the selected box\n" \ "Ctrl + R\t\tRe-recognize the selected box\n" \ "Ctrl + C\t\tCopy and paste the selected\n" \ "\t\t\tbox\n" \ - "\n"\ + "\n" \ "Ctrl + Left Mouse\tMulti select the label\n" \ "Button\t\t\tbox\n" \ - "\n"\ + "\n" \ "Backspace\t\tDelete the selected box\n" \ "Ctrl + V\t\tCheck image\n" \ "Ctrl + Shift + d\tDelete image\n" \ @@ -245,4 +244,4 @@ def keysInfo(lang='en'): "———————————————————————\n" \ "Notice:For Mac users, use the 'Command' key instead of the 'Ctrl' key" - return msg \ No newline at end of file + return msg diff --git a/PPOCRLabel/resources/strings/strings-en.properties b/PPOCRLabel/resources/strings/strings-en.properties index f59e43aa92ff9ccd04686e9c16db181983b57b2c..3c4eda65a32e1048405041667ba61bdb639bfd7b 100644 --- a/PPOCRLabel/resources/strings/strings-en.properties +++ b/PPOCRLabel/resources/strings/strings-en.properties @@ -106,4 +106,7 @@ undo=Undo undoLastPoint=Undo Last Point autoSaveMode=Auto Export Label Mode lockBox=Lock selected box/Unlock all box -lockBoxDetail=Lock selected box/Unlock all box \ No newline at end of file +lockBoxDetail=Lock selected box/Unlock all box +keyListTitle=Key List +keyDialogTip=Enter object label +keyChange=Change Box Key diff --git a/PPOCRLabel/resources/strings/strings-zh-CN.properties b/PPOCRLabel/resources/strings/strings-zh-CN.properties index d8bd9d4bff02748397d7a57a6205e67ff69779c2..a7c30368b87354cbae81b2cdead8ad31b2a8c1eb 100644 --- a/PPOCRLabel/resources/strings/strings-zh-CN.properties +++ b/PPOCRLabel/resources/strings/strings-zh-CN.properties @@ -107,3 +107,6 @@ undoLastPoint=撤销上个点 autoSaveMode=自动导出标记结果 lockBox=锁定框/解除锁定框 lockBoxDetail=若当前没有框处于锁定状态则锁定选中的框,若存在锁定框则解除所有锁定框的锁定状态 +keyListTitle=关键词列表 +keyDialogTip=请输入类型名称 +keyChange=更改Box关键字类别 \ No newline at end of file