diff --git a/PPOCRLabel/PPOCRLabel.py b/PPOCRLabel/PPOCRLabel.py index fe41746060a4f82913326011bb7b4674043677f8..4d6c0d7770cd990abd26f15285d2c6542fc54fc7 100644 --- a/PPOCRLabel/PPOCRLabel.py +++ b/PPOCRLabel/PPOCRLabel.py @@ -202,11 +202,8 @@ class MainWindow(QMainWindow): self.reRecogButton.setIcon(newIcon('reRec', 30)) self.reRecogButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon) - self.cellreRecButton = QToolButton() - self.cellreRecButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon) self.tableRecButton = QToolButton() self.tableRecButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon) - self.newButton = QToolButton() self.newButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon) @@ -221,9 +218,9 @@ class MainWindow(QMainWindow): leftTopToolBox = QGridLayout() leftTopToolBox.addWidget(self.newButton, 0, 0, 1, 1) leftTopToolBox.addWidget(self.createpolyButton, 0, 1, 1, 1) - leftTopToolBox.addWidget(self.reRecogButton, 0, 2, 1, 1) - leftTopToolBox.addWidget(self.tableRecButton, 1, 0, 1, 1) - leftTopToolBox.addWidget(self.cellreRecButton, 1, 1, 1, 1) + leftTopToolBox.addWidget(self.reRecogButton, 1, 0, 1, 1) + leftTopToolBox.addWidget(self.tableRecButton, 1, 1, 1, 1) + leftTopToolBoxContainer = QWidget() leftTopToolBoxContainer.setLayout(leftTopToolBox) listLayout.addWidget(leftTopToolBoxContainer) @@ -507,7 +504,6 @@ class MainWindow(QMainWindow): self.AutoRecognition.setDefaultAction(AutoRec) self.reRecogButton.setDefaultAction(reRec) self.tableRecButton.setDefaultAction(tableRec) - self.cellreRecButton.setDefaultAction(cellreRec) # self.preButton.setDefaultAction(openPrevImg) # self.nextButton.setDefaultAction(openNextImg) @@ -564,11 +560,11 @@ class MainWindow(QMainWindow): rotateLeft=rotateLeft, rotateRight=rotateRight, lock=lock, exportJSON=exportJSON, fileMenuActions=(opendir, open_dataset_dir, saveLabel, exportJSON, resetAll, quit), beginner=(), advanced=(), - editMenu=(createpoly, edit, copy, delete, singleRere, None, undo, undoLastPoint, + editMenu=(createpoly, edit, copy, delete, singleRere, cellreRec, None, undo, undoLastPoint, None, rotateLeft, rotateRight, None, color1, self.drawSquaresOption, lock, None, change_cls), beginnerContext=( - create, createpoly, edit, copy, delete, singleRere, rotateLeft, rotateRight, lock, change_cls), + create, createpoly, edit, copy, delete, singleRere, cellreRec, rotateLeft, rotateRight, lock, change_cls), advancedContext=(createMode, editMode, edit, copy, delete, shapeLineColor, shapeFillColor), onLoadActive=(create, createpoly, createMode, editMode), @@ -1025,6 +1021,7 @@ class MainWindow(QMainWindow): self._noSelectionSlot = False n_selected = len(selected_shapes) self.actions.singleRere.setEnabled(n_selected) + self.actions.cellreRec.setEnabled(n_selected) self.actions.delete.setEnabled(n_selected) self.actions.copy.setEnabled(n_selected) self.actions.edit.setEnabled(n_selected == 1) @@ -1690,12 +1687,10 @@ class MainWindow(QMainWindow): self.haveAutoReced = False self.AutoRecognition.setEnabled(True) self.reRecogButton.setEnabled(True) - self.cellreRecButton.setEnabled(True) self.tableRecButton.setEnabled(True) self.actions.AutoRec.setEnabled(True) self.actions.reRec.setEnabled(True) self.actions.tableRec.setEnabled(True) - self.actions.cellreRec.setEnabled(True) self.actions.open_dataset_dir.setEnabled(True) self.actions.rotateLeft.setEnabled(True) self.actions.rotateRight.setEnabled(True) @@ -2229,87 +2224,59 @@ class MainWindow(QMainWindow): re-recognise text in a cell ''' img = cv2.imread(self.filePath) + for shape in self.canvas.selectedShapes: + box = [[int(p.x()), int(p.y())] for p in shape.points] - if self.canvas.shapes: - self.result_dic = [] - self.result_dic_locked = [] # result_dic_locked stores the ocr result of self.canvas.lockedShapes - rec_flag = 0 - for shape in self.canvas.shapes: - box = [[int(p.x()), int(p.y())] for p in shape.points] + if len(box) > 4: + box = self.gen_quad_from_poly(np.array(box)) + assert len(box) == 4 - if len(box) > 4: - box = self.gen_quad_from_poly(np.array(box)) - assert len(box) == 4 + # pad around bbox for better text recognition accuracy + _box = boxPad(box, img.shape, 6) + 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 - # pad around bbox for better text recognition accuracy - print(box) - _box = boxPad(box, img.shape, 6) - print(_box) - 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 - # merge the text result in the cell - texts = '' - probs = 0. # the probability of the cell is avgerage prob of every text box in the cell - bboxes = self.ocr.ocr(img_crop, det=True, rec=False, cls=False) - if len(bboxes) > 0: - bboxes.reverse() # top row text at first - for _bbox in bboxes: - patch = get_rotate_crop_image(img_crop, np.array(_bbox, np.float32)) - rec_res = self.ocr.ocr(patch, det=False, rec=True, cls=False) - text = rec_res[0][0] - if text != '': - texts += text + (' ' if text[0].isalpha() else '') # add space between english word - probs += rec_res[0][1] - probs = probs / len(bboxes) - result = [(texts.strip(), probs)] + # merge the text result in the cell + texts = '' + probs = 0. # the probability of the cell is avgerage prob of every text box in the cell + bboxes = self.ocr.ocr(img_crop, det=True, rec=False, cls=False) + if len(bboxes) > 0: + bboxes.reverse() # top row text at first + for _bbox in bboxes: + patch = get_rotate_crop_image(img_crop, np.array(_bbox, np.float32)) + rec_res = self.ocr.ocr(patch, det=False, rec=True, cls=False) + text = rec_res[0][0] + if text != '': + texts += text + (' ' if text[0].isalpha() else '') # add space between english word + probs += rec_res[0][1] + probs = probs / len(bboxes) + result = [(texts.strip(), probs)] - if result[0][0] != '': - if shape.line_color == DEFAULT_LOCK_COLOR: - shape.label = result[0][0] - result.insert(0, box) - self.result_dic_locked.append(result) - else: - result.insert(0, box) - self.result_dic.append(result) - else: - print('Can not recognise the box') - if shape.line_color == DEFAULT_LOCK_COLOR: - shape.label = result[0][0] - self.result_dic_locked.append([box, (self.noLabelText, 0)]) - else: - self.result_dic.append([box, (self.noLabelText, 0)]) - try: - if self.noLabelText == shape.label or result[1][0] == shape.label: - print('label no change') - else: - rec_flag += 1 - except IndexError as e: - print('Can not recognise the box') - if (len(self.result_dic) > 0 and rec_flag > 0) or self.canvas.lockedShapes: - self.canvas.isInTheSameImage = True - self.saveFile(mode='Auto') - self.loadFile(self.filePath) - self.canvas.isInTheSameImage = False - self.setDirty() - elif len(self.result_dic) == len(self.canvas.shapes) and rec_flag == 0: - if self.lang == 'ch': - QMessageBox.information(self, "Information", "识别结果保持一致！") + 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: - QMessageBox.information(self, "Information", "The recognition result remains unchanged!") + shape.label = result[1][0] else: - print('Can not recgonise in ', self.filePath) - else: - QMessageBox.information(self, "Information", "Draw a box!") + print('Can not recognise the box') + if self.noLabelText == shape.label: + print('label no change') + else: + shape.label = self.noLabelText + self.singleLabel(shape) + self.setDirty() def exportJSON(self): ''' export PPLabel and CSV to JSON (PubTabNet) ''' import pandas as pd - from PyQt5.QtWidgets import QInputDialog + from libs.dataPartitionDialog import DataPartitionDialog if self.lang == 'ch': QMessageBox.information(self, "Information", "导出JSON前请保存所有图像的标注且关闭EXCEL!!!!!!!!!!!!") @@ -2346,18 +2313,18 @@ class MainWindow(QMainWindow): # return # data partition user input - train_split, ok = QInputDialog.getInt(self, "DataPatition", "How many data for Training (%):", 70, 0, 100, 1) - if not ok: - return - val_split, ok = QInputDialog.getInt(self, "DataPatition", "How many data for Validatiion (%):", 15, 0, 100, 1) - if not ok: + partitionDialog = DataPartitionDialog() + partitionDialog.exec() + if partitionDialog.getStatus() == False: return - test_split, ok = QInputDialog.getInt(self, "DataPatition", "How many data for Testing (%):", 15, 0, 100, 1) + train_split, val_split, test_split = partitionDialog.getDataPartition() + # check validate if train_split + val_split + test_split > 100: - QMessageBox.information(self, "Information", "The sum of training, validation and testing data should be less than 100%") + msg = "The sum of training, validation and testing data should be less than 100%" + QMessageBox.information(self, "Information", msg) return - + print(train_split, val_split, test_split) train_split, val_split, test_split = float(train_split) / 100., float(val_split) / 100., float(test_split) / 100. train_id = int(len(labeldict) * train_split) val_id = int(len(labeldict) * (train_split + val_split)) @@ -2407,7 +2374,7 @@ class MainWindow(QMainWindow): with open("{}/annotation.json".format(self.lastOpenDir), "w") as fid: fid.write(json.dumps(json_results)) - msg = 'JSON sucessfully saved in ', "{}/annotation.json".format(self.lastOpenDir) + msg = 'JSON sucessfully saved in {}/annotation.json'.format(self.lastOpenDir) QMessageBox.information(self, "Information", msg) def autolcm(self): diff --git a/PPOCRLabel/libs/dataPartitionDialog.py b/PPOCRLabel/libs/dataPartitionDialog.py new file mode 100644 index 0000000000000000000000000000000000000000..051b9722d56bbd6c0bbf61f88e34c9b57af9ddad --- /dev/null +++ b/PPOCRLabel/libs/dataPartitionDialog.py @@ -0,0 +1,101 @@ +try: + from PyQt5.QtGui import * + from PyQt5.QtCore import * + from PyQt5.QtWidgets import * +except ImportError: + from PyQt4.QtGui import * + from PyQt4.QtCore import * + +from libs.utils import newIcon + +import time +import datetime +import json +import cv2 +import numpy as np + + +BB = QDialogButtonBox + +class DataPartitionDialog(QDialog): + def __init__(self): + super().__init__() + self.title = 'DATA PARTITION' + + self.train_ratio = 70 + self.val_ratio = 15 + self.test_ratio = 15 + + self.initUI() + + def initUI(self): + self.setWindowTitle(self.title) + self.setWindowModality(Qt.ApplicationModal) + + self.flag_accept = True + + train_lbl = QLabel('Train split: ', self) + train_lbl.setFont(QFont('Arial', 15)) + val_lbl = QLabel('Valid split: ', self) + val_lbl.setFont(QFont('Arial', 15)) + test_lbl = QLabel('Test split: ', self) + test_lbl.setFont(QFont('Arial', 15)) + + self.train_input = QLineEdit(self) + self.train_input.setFont(QFont('Arial', 15)) + self.val_input = QLineEdit(self) + self.val_input.setFont(QFont('Arial', 15)) + self.test_input = QLineEdit(self) + self.test_input.setFont(QFont('Arial', 15)) + + self.train_input.setText(str(self.train_ratio)) + self.val_input.setText(str(self.val_ratio)) + self.test_input.setText(str(self.test_ratio)) + + validator = QIntValidator(0, 100) + self.train_input.setValidator(validator) + self.val_input.setValidator(validator) + self.test_input.setValidator(validator) + + gridlayout = QGridLayout() + gridlayout.addWidget(train_lbl, 0, 0) + gridlayout.addWidget(val_lbl, 1, 0) + gridlayout.addWidget(test_lbl, 2, 0) + gridlayout.addWidget(self.train_input, 0, 1) + gridlayout.addWidget(self.val_input, 1, 1) + gridlayout.addWidget(self.test_input, 2, 1) + + bb = BB(BB.Ok | BB.Cancel, 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.cancel) + gridlayout.addWidget(bb, 3, 0, 1, 2) + + self.setLayout(gridlayout) + + self.show() + + def validate(self): + self.flag_accept = True + self.accept() + + def cancel(self): + self.flag_accept = False + self.reject() + + def getStatus(self): + return self.flag_accept + + def getDataPartition(self): + self.train_ratio = int(self.train_input.text()) + self.val_ratio = int(self.val_input.text()) + self.test_ratio = int(self.test_input.text()) + + return self.train_ratio, self.val_ratio, self.test_ratio + + def closeEvent(self, event): + self.flag_accept = False + self.reject() + + diff --git a/PPOCRLabel/resources/strings/strings-en.properties b/PPOCRLabel/resources/strings/strings-en.properties index f2e289ff75daa5975606e3b02000eaaf600f2f1a..7ba9af4c33ebe8130b5ce529e01270361c5f11d6 100644 --- a/PPOCRLabel/resources/strings/strings-en.properties +++ b/PPOCRLabel/resources/strings/strings-en.properties @@ -84,7 +84,7 @@ mhelp=Help iconList=Icon List detectionBoxposition=Detection box position recognitionResult=Recognition result -creatPolygon=Create Quadrilateral +creatPolygon=Create PolygonBox rotateLeft=Left turn 90 degrees rotateRight=Right turn 90 degrees drawSquares=Draw Squares diff --git a/PPOCRLabel/resources/strings/strings-zh-CN.properties b/PPOCRLabel/resources/strings/strings-zh-CN.properties index 7cc0ea2286618d62e729030822654c90462d77ca..308974ef29fac367bccd5f0e97926b2087a9f8df 100644 --- a/PPOCRLabel/resources/strings/strings-zh-CN.properties +++ b/PPOCRLabel/resources/strings/strings-zh-CN.properties @@ -84,7 +84,7 @@ mhelp=帮助 iconList=缩略图 detectionBoxposition=检测框位置 recognitionResult=识别结果 -creatPolygon=四点标注 +creatPolygon=多边形标注 drawSquares=正方形标注 rotateLeft=图片左旋转90度 rotateRight=图片右旋转90度