From 320cb381064e723aeb76f1c0d04cda1420515e10 Mon Sep 17 00:00:00 2001 From: whj_dark Date: Wed, 4 May 2022 13:59:44 +0800 Subject: [PATCH 1/7] new --- PPOCRLabel/PPOCRLabel.py | 364 +++++++++++++++++- PPOCRLabel/libs/utils.py | 27 ++ .../resources/strings/strings-en.properties | 3 + .../strings/strings-zh-CN.properties | 5 +- 4 files changed, 383 insertions(+), 16 deletions(-) diff --git a/PPOCRLabel/PPOCRLabel.py b/PPOCRLabel/PPOCRLabel.py index 523661b635..fe41746060 100644 --- a/PPOCRLabel/PPOCRLabel.py +++ b/PPOCRLabel/PPOCRLabel.py @@ -26,7 +26,7 @@ from functools import partial 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, \ + QSlider, QGraphicsOpacityEffect, QMessageBox, QListView, QScrollArea, QWidgetAction, QApplication, QLabel, QGridLayout, \ QFileDialog, QListWidgetItem, QComboBox, QDialog __dir__ = os.path.dirname(os.path.abspath(__file__)) @@ -36,7 +36,7 @@ sys.path.append(os.path.abspath(os.path.join(__dir__, '../..'))) sys.path.append(os.path.abspath(os.path.join(__dir__, '../PaddleOCR'))) sys.path.append("..") -from paddleocr import PaddleOCR +from paddleocr import PaddleOCR, PPStructure from libs.constants import * from libs.utils import * from libs.labelColor import label_colormap @@ -100,9 +100,15 @@ class MainWindow(QMainWindow): use_gpu=gpu, lang=lang, show_log=False) + self.table_ocr = PPStructure(use_pdserving=False, + use_gpu=gpu, + lang=lang, + layout=False, + show_log=False) if os.path.exists('./data/paddle.png'): result = self.ocr.ocr('./data/paddle.png', cls=True, det=True) + result = self.table_ocr('./data/paddle.png', return_ocr_result_in_table=True) # For loading all image under a directory self.mImgList = [] @@ -196,16 +202,28 @@ 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) + self.createpolyButton = QToolButton() + self.createpolyButton.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) + 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) leftTopToolBoxContainer = QWidget() leftTopToolBoxContainer.setLayout(leftTopToolBox) listLayout.addWidget(leftTopToolBoxContainer) @@ -446,13 +464,22 @@ class MainWindow(QMainWindow): 'Ctrl+R', 'reRec', getStr('singleRe'), enabled=False) createpoly = action(getStr('creatPolygon'), self.createPolygon, - 'q', 'new', getStr('creatPolygon'), enabled=True) + 'q', 'new', getStr('creatPolygon'), enabled=False) + + tableRec = action(getStr('TableRecognition'), self.TableRecognition, + '', 'Auto', getStr('TableRecognition'), enabled=False) + + cellreRec = action(getStr('cellreRecognition'), self.cellreRecognition, + '', 'reRec', getStr('cellreRecognition'), enabled=False) saveRec = action(getStr('saveRec'), self.saveRecResult, '', 'save', getStr('saveRec'), enabled=False) saveLabel = action(getStr('saveLabel'), self.saveLabelFile, # 'Ctrl+S', 'save', getStr('saveLabel'), enabled=False) + + exportJSON = action(getStr('exportJSON'), self.exportJSON, + '', 'save', getStr('exportJSON'), enabled=False) undoLastPoint = action(getStr("undoLastPoint"), self.canvas.undoLastPoint, 'Ctrl+Z', "undo", getStr("undoLastPoint"), enabled=False) @@ -474,10 +501,13 @@ class MainWindow(QMainWindow): self.editButton.setDefaultAction(edit) self.newButton.setDefaultAction(create) + self.createpolyButton.setDefaultAction(createpoly) self.DelButton.setDefaultAction(deleteImg) self.SaveButton.setDefaultAction(save) self.AutoRecognition.setDefaultAction(AutoRec) self.reRecogButton.setDefaultAction(reRec) + self.tableRecButton.setDefaultAction(tableRec) + self.cellreRecButton.setDefaultAction(cellreRec) # self.preButton.setDefaultAction(openPrevImg) # self.nextButton.setDefaultAction(openNextImg) @@ -523,25 +553,25 @@ class MainWindow(QMainWindow): # 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, + lineColor=color1, create=create, createpoly=createpoly, tableRec=tableRec, delete=delete, edit=edit, copy=copy, + saveRec=saveRec, singleRere=singleRere, AutoRec=AutoRec, reRec=reRec, cellreRec=cellreRec, createMode=createMode, editMode=editMode, shapeLineColor=shapeLineColor, shapeFillColor=shapeFillColor, zoom=zoom, zoomIn=zoomIn, zoomOut=zoomOut, zoomOrg=zoomOrg, fitWindow=fitWindow, fitWidth=fitWidth, 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), + 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, None, rotateLeft, rotateRight, None, color1, self.drawSquaresOption, lock, None, change_cls), beginnerContext=( - create, edit, copy, delete, singleRere, rotateLeft, rotateRight, lock, change_cls), + create, createpoly, edit, copy, delete, singleRere, rotateLeft, rotateRight, lock, change_cls), advancedContext=(createMode, editMode, edit, copy, delete, shapeLineColor, shapeFillColor), - onLoadActive=(create, createMode, editMode), + onLoadActive=(create, createpoly, createMode, editMode), onShapesPresent=(hideAll, showAll)) # menus @@ -574,7 +604,7 @@ class MainWindow(QMainWindow): self.autoSaveOption.triggered.connect(self.autoSaveFunc) addActions(self.menus.file, - (opendir, open_dataset_dir, None, saveLabel, saveRec, self.autoSaveOption, None, resetAll, deleteImg, + (opendir, open_dataset_dir, None, saveLabel, saveRec, exportJSON, self.autoSaveOption, None, resetAll, deleteImg, quit)) addActions(self.menus.help, (showKeys, showSteps, showInfo)) @@ -585,7 +615,7 @@ class MainWindow(QMainWindow): zoomIn, zoomOut, zoomOrg, None, fitWindow, fitWidth)) - addActions(self.menus.autolabel, (AutoRec, reRec, alcm, None, help)) + addActions(self.menus.autolabel, (AutoRec, reRec, cellreRec, alcm, None, help)) self.menus.file.aboutToShow.connect(self.updateFileMenu) @@ -695,6 +725,7 @@ class MainWindow(QMainWindow): self.dirty = False self.actions.save.setEnabled(False) self.actions.create.setEnabled(True) + self.actions.createpoly.setEnabled(True) def toggleActions(self, value=True): """Enable/Disable widgets which depend on an opened image.""" @@ -780,6 +811,7 @@ class MainWindow(QMainWindow): assert self.beginner() self.canvas.setEditing(False) self.actions.create.setEnabled(False) + self.actions.createpoly.setEnabled(False) self.canvas.fourpoint = False def createPolygon(self): @@ -787,10 +819,10 @@ class MainWindow(QMainWindow): self.canvas.setEditing(False) self.canvas.fourpoint = True self.actions.create.setEnabled(False) + self.actions.createpoly.setEnabled(False) self.actions.undoLastPoint.setEnabled(True) def rotateImg(self, filename, k, _value): - self.actions.rotateRight.setEnabled(_value) pix = cv2.imread(filename) pix = np.rot90(pix, k) @@ -831,6 +863,7 @@ class MainWindow(QMainWindow): self.canvas.setEditing(True) self.canvas.restoreCursor() self.actions.create.setEnabled(True) + self.actions.createpoly.setEnabled(True) def toggleDrawMode(self, edit=True): self.canvas.setEditing(edit) @@ -1216,6 +1249,7 @@ class MainWindow(QMainWindow): if self.beginner(): # Switch to edit mode. self.canvas.setEditing(True) self.actions.create.setEnabled(True) + self.actions.createpoly.setEnabled(True) self.actions.undoLastPoint.setEnabled(False) self.actions.undo.setEnabled(True) else: @@ -1656,8 +1690,12 @@ 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) @@ -1757,6 +1795,7 @@ class MainWindow(QMainWindow): self.openNextImg() self.actions.saveRec.setEnabled(True) self.actions.saveLabel.setEnabled(True) + self.actions.exportJSON.setEnabled(True) elif mode == 'Auto': if annotationFilePath and self.saveLabels(annotationFilePath, mode=mode): @@ -2083,6 +2122,294 @@ class MainWindow(QMainWindow): self.singleLabel(shape) self.setDirty() + def TableRecognition(self): + ''' + Table Recegnition + ''' + from paddleocr.ppstructure.table.predict_table import to_excel + + import time + + start = time.time() + img = cv2.imread(self.filePath) + res = self.table_ocr(img, return_ocr_result_in_table=True) + + TableRec_excel_dir = self.lastOpenDir + '/tableRec_excel_output/' + os.makedirs(TableRec_excel_dir, exist_ok=True) + filename = os.path.basename(self.filePath) + excel_path = TableRec_excel_dir + '{}.xlsx'.format(filename) + + if res is None: + msg = 'Can not recognise the table in ' + self.filePath + '. Please change manually' + QMessageBox.information(self, "Information", msg) + to_excel('', excel_path) # create an empty excel + return + + # save res + # ONLY SUPPORT ONE TABLE in one image + hasTable = False + for region in res: + if region['type'] == 'Table': + if region['res']['boxes'] is None: + msg = 'Can not recognise the detection box in ' + self.filePath + '. Please change manually' + QMessageBox.information(self, "Information", msg) + to_excel('', excel_path) # create an empty excel + return + hasTable = True + # save table ocr result on PPOCRLabel + # clear all old annotaions before saving result + self.itemsToShapes.clear() + self.shapesToItems.clear() + self.itemsToShapesbox.clear() # ADD + self.shapesToItemsbox.clear() + self.labelList.clear() + self.BoxList.clear() + self.result_dic = [] + self.result_dic_locked = [] + + shapes = [] + result_len = len(region['res']['boxes']) + for i in range(result_len): + bbox = np.array(region['res']['boxes'][i]) + rec_text = region['res']['rec_res'][i][0] + + # polys to rectangles + x1, y1 = np.min(bbox[:, 0]), np.min(bbox[:, 1]) + x2, y2 = np.max(bbox[:, 0]), np.max(bbox[:, 1]) + rext_bbox = [[x1, y1], [x2, y1], [x2, y2], [x1, y2]] + + # save bbox to shape + shape = Shape(label=rec_text, line_color=DEFAULT_LINE_COLOR, key_cls=None) + for point in rext_bbox: + x, y = point + # Ensure the labels are within the bounds of the image. + # If not, fix them. + x, y, snapped = self.canvas.snapPointToCanvas(x, y) + shape.addPoint(QPointF(x, y)) + shape.difficult = False + # shape.locked = False + shape.close() + self.addLabel(shape) + shapes.append(shape) + self.setDirty() + self.canvas.loadShapes(shapes) + + # save HTML result to excel + try: + to_excel(region['res']['html'], excel_path) + except: + print('Can not save excel file, maybe Permission denied (.xlsx is being occupied)') + break + + if not hasTable: + msg = 'Can not recognise the table in ' + self.filePath + '. Please change manually' + QMessageBox.information(self, "Information", msg) + to_excel('', excel_path) # create an empty excel + return + + # automatically open excel annotation file + try: + import win32com.client + except: + print("CANNOT OPEN .xlsx. It could be one of the following reasons: " \ + "Only support Windows | No python win32com") + + try: + xl = win32com.client.Dispatch("Excel.Application") + xl.Visible = True + xl.Workbooks.Open(excel_path) + except: + print("CANNOT OPEN .xlsx. It could be the following reasons: " \ + ".xlsx is not existed") + + print('time cost: ', time.time() - start) + + def cellreRecognition(self): + ''' + re-recognise text in a cell + ''' + img = cv2.imread(self.filePath) + + 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 + + # 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)] + + 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", "识别结果保持一致！") + else: + 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 exportJSON(self): + ''' + export PPLabel and CSV to JSON (PubTabNet) + ''' + import pandas as pd + from PyQt5.QtWidgets import QInputDialog + + if self.lang == 'ch': + QMessageBox.information(self, "Information", "导出JSON前请保存所有图像的标注且关闭EXCEL!!!!!!!!!!!!") + else: + QMessageBox.information(self, "Information", "Please save all the annotations and close the EXCEL before exporting JSON!!!!!!!!!!!!") + + # automatically save annotations + self.saveLabelFile() + + # load box annotations + labeldict = {} + if not os.path.exists(self.PPlabelpath): + msg = 'ERROR, Can not find Label.txt' + QMessageBox.information(self, "Information", msg) + return + else: + with open(self.PPlabelpath, '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] = [] + + # if len(labeldict) != len(csv_paths): + # msg = 'ERROR, box label and excel label are not in the same number\n' + \ + # 'box label: ' + str(len(labeldict)) + '\n' + \ + # 'excel label: ' + str(len(csv_paths)) + '\n' + \ + # 'Please check the label.txt and tableRec_excel_output\n' + # QMessageBox.information(self, "Information", msg) + # 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: + return + test_split, ok = QInputDialog.getInt(self, "DataPatition", "How many data for Testing (%):", 15, 0, 100, 1) + + 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%") + return + + 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)) + print('Data partition: train:', train_id, + 'validation:', val_id - train_id, + 'test:', len(labeldict) - val_id) + + TableRec_excel_dir = os.path.join(self.lastOpenDir, 'tableRec_excel_output') + json_results = [] + imgid = 0 + for image_path in labeldict.keys(): + # load csv annotations + filename = os.path.basename(image_path) + csv_path = os.path.join(TableRec_excel_dir, filename + '.xlsx') + if not os.path.exists(csv_path): + msg = 'ERROR, Can not find ' + csv_path + QMessageBox.information(self, "Information", msg) + return + + # read xlsx file, convert to HTML + xd = pd.ExcelFile(csv_path) + df = xd.parse() + structure = df.to_html() + + # load box annotations + cells = [] + for anno in labeldict[image_path]: + tokens = list(anno['transcription']) + obb = anno['points'] + hbb = OBB2HBB(np.array(obb)).tolist() + cells.append({'tokens': tokens, 'bbox': hbb}) + + # data split + if imgid < train_id: + split = 'train' + elif imgid < val_id: + split = 'val' + else: + split = 'test' + + # save dict + html = {'structure': {'tokens': structure}, 'cell': cells} + json_results.append({'filename': filename, 'split': split, 'imgid': imgid, 'html': html}) + imgid += 1 + + # save json + 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) + QMessageBox.information(self, "Information", msg) + def autolcm(self): vbox = QVBoxLayout() hbox = QHBoxLayout() @@ -2122,6 +2449,12 @@ class MainWindow(QMainWindow): 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()]) + del self.table_ocr + self.table_ocr = PPStructure(use_pdserving=False, + use_gpu=False, + lang=lg_idx[self.comboBox.currentText()], + layout=False, + show_log=False) self.dialog.close() def cancel(self): @@ -2140,6 +2473,7 @@ class MainWindow(QMainWindow): self.fileStatedict[file] = 1 self.actions.saveLabel.setEnabled(True) self.actions.saveRec.setEnabled(True) + self.actions.exportJSON.setEnabled(True) def saveFilestate(self): with open(self.fileStatepath, 'w', encoding='utf-8') as f: diff --git a/PPOCRLabel/libs/utils.py b/PPOCRLabel/libs/utils.py index 2510520caa..c49b506882 100644 --- a/PPOCRLabel/libs/utils.py +++ b/PPOCRLabel/libs/utils.py @@ -161,6 +161,33 @@ def get_rotate_crop_image(img, points): print(e) +def boxPad(box, imgShape, pad : int) -> np.array: + """ + Pad a box with [pad] pixels on each side. + """ + box = np.array(box, dtype=np.int32) + box[0][0], box[0][1] = box[0][0] - pad, box[0][1] - pad + box[1][0], box[1][1] = box[1][0] + pad, box[1][1] - pad + box[2][0], box[2][1] = box[2][0] + pad, box[2][1] + pad + box[3][0], box[3][1] = box[3][0] - pad, box[3][1] + pad + h, w, _ = imgShape + box[:,0] = np.clip(box[:,0], 0, w) + box[:,1] = np.clip(box[:,1], 0, h) + return box + + +def OBB2HBB(obb) -> np.array: + """ + Convert Oriented Bounding Box to Horizontal Bounding Box. + """ + hbb = np.zeros(4, dtype=np.int32) + hbb[0] = min(obb[:, 0]) + hbb[1] = min(obb[:, 1]) + hbb[2] = max(obb[:, 0]) + hbb[3] = max(obb[:, 1]) + return hbb + + def stepsInfo(lang='en'): if lang == 'ch': msg = "1. 安装与运行：使用上述命令安装与运行程序。\n" \ diff --git a/PPOCRLabel/resources/strings/strings-en.properties b/PPOCRLabel/resources/strings/strings-en.properties index 3c4eda65a3..f2e289ff75 100644 --- a/PPOCRLabel/resources/strings/strings-en.properties +++ b/PPOCRLabel/resources/strings/strings-en.properties @@ -110,3 +110,6 @@ lockBoxDetail=Lock selected box/Unlock all box keyListTitle=Key List keyDialogTip=Enter object label keyChange=Change Box Key +TableRecognition=Table Recognition +cellreRecognition=Cell Re-Recognition +exportJSON=export JSON(PubTabNet) diff --git a/PPOCRLabel/resources/strings/strings-zh-CN.properties b/PPOCRLabel/resources/strings/strings-zh-CN.properties index a7c30368b8..7cc0ea2286 100644 --- a/PPOCRLabel/resources/strings/strings-zh-CN.properties +++ b/PPOCRLabel/resources/strings/strings-zh-CN.properties @@ -109,4 +109,7 @@ lockBox=锁定框/解除锁定框 lockBoxDetail=若当前没有框处于锁定状态则锁定选中的框，若存在锁定框则解除所有锁定框的锁定状态 keyListTitle=关键词列表 keyDialogTip=请输入类型名称 -keyChange=更改Box关键字类别 \ No newline at end of file +keyChange=更改Box关键字类别 +TableRecognition=表格识别 +cellreRecognition=单元格重识别 +exportJSON=导出表格JSON标注 \ No newline at end of file From 8b228a1f9b011aba935963431cadb81c7fe361d5 Mon Sep 17 00:00:00 2001 From: whj_dark Date: Wed, 4 May 2022 13:59:44 +0800 Subject: [PATCH 2/7] new --- PPOCRLabel/PPOCRLabel.py | 331 +++++++++++++++++- PPOCRLabel/libs/dataPartitionDialog.py | 101 ++++++ PPOCRLabel/libs/utils.py | 27 ++ .../resources/strings/strings-en.properties | 5 +- .../strings/strings-zh-CN.properties | 7 +- 5 files changed, 453 insertions(+), 18 deletions(-) create mode 100644 PPOCRLabel/libs/dataPartitionDialog.py diff --git a/PPOCRLabel/PPOCRLabel.py b/PPOCRLabel/PPOCRLabel.py index 523661b635..7fc174e09f 100644 --- a/PPOCRLabel/PPOCRLabel.py +++ b/PPOCRLabel/PPOCRLabel.py @@ -26,7 +26,7 @@ from functools import partial 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, \ + QSlider, QGraphicsOpacityEffect, QMessageBox, QListView, QScrollArea, QWidgetAction, QApplication, QLabel, QGridLayout, \ QFileDialog, QListWidgetItem, QComboBox, QDialog __dir__ = os.path.dirname(os.path.abspath(__file__)) @@ -36,7 +36,7 @@ sys.path.append(os.path.abspath(os.path.join(__dir__, '../..'))) sys.path.append(os.path.abspath(os.path.join(__dir__, '../PaddleOCR'))) sys.path.append("..") -from paddleocr import PaddleOCR +from paddleocr import PaddleOCR, PPStructure from libs.constants import * from libs.utils import * from libs.labelColor import label_colormap @@ -100,9 +100,15 @@ class MainWindow(QMainWindow): use_gpu=gpu, lang=lang, show_log=False) + self.table_ocr = PPStructure(use_pdserving=False, + use_gpu=gpu, + lang=lang, + layout=False, + show_log=False) if os.path.exists('./data/paddle.png'): result = self.ocr.ocr('./data/paddle.png', cls=True, det=True) + result = self.table_ocr('./data/paddle.png', return_ocr_result_in_table=True) # For loading all image under a directory self.mImgList = [] @@ -196,16 +202,25 @@ class MainWindow(QMainWindow): self.reRecogButton.setIcon(newIcon('reRec', 30)) self.reRecogButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon) + self.tableRecButton = QToolButton() + self.tableRecButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon) + self.newButton = QToolButton() self.newButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon) + self.createpolyButton = QToolButton() + self.createpolyButton.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) + leftTopToolBox = QGridLayout() + leftTopToolBox.addWidget(self.newButton, 0, 0, 1, 1) + leftTopToolBox.addWidget(self.createpolyButton, 0, 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) @@ -446,13 +461,22 @@ class MainWindow(QMainWindow): 'Ctrl+R', 'reRec', getStr('singleRe'), enabled=False) createpoly = action(getStr('creatPolygon'), self.createPolygon, - 'q', 'new', getStr('creatPolygon'), enabled=True) + 'q', 'new', getStr('creatPolygon'), enabled=False) + + tableRec = action(getStr('TableRecognition'), self.TableRecognition, + '', 'Auto', getStr('TableRecognition'), enabled=False) + + cellreRec = action(getStr('cellreRecognition'), self.cellreRecognition, + '', 'reRec', getStr('cellreRecognition'), enabled=False) saveRec = action(getStr('saveRec'), self.saveRecResult, '', 'save', getStr('saveRec'), enabled=False) saveLabel = action(getStr('saveLabel'), self.saveLabelFile, # 'Ctrl+S', 'save', getStr('saveLabel'), enabled=False) + + exportJSON = action(getStr('exportJSON'), self.exportJSON, + '', 'save', getStr('exportJSON'), enabled=False) undoLastPoint = action(getStr("undoLastPoint"), self.canvas.undoLastPoint, 'Ctrl+Z', "undo", getStr("undoLastPoint"), enabled=False) @@ -474,10 +498,12 @@ class MainWindow(QMainWindow): self.editButton.setDefaultAction(edit) self.newButton.setDefaultAction(create) + self.createpolyButton.setDefaultAction(createpoly) self.DelButton.setDefaultAction(deleteImg) self.SaveButton.setDefaultAction(save) self.AutoRecognition.setDefaultAction(AutoRec) self.reRecogButton.setDefaultAction(reRec) + self.tableRecButton.setDefaultAction(tableRec) # self.preButton.setDefaultAction(openPrevImg) # self.nextButton.setDefaultAction(openNextImg) @@ -523,25 +549,25 @@ class MainWindow(QMainWindow): # 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, + lineColor=color1, create=create, createpoly=createpoly, tableRec=tableRec, delete=delete, edit=edit, copy=copy, + saveRec=saveRec, singleRere=singleRere, AutoRec=AutoRec, reRec=reRec, cellreRec=cellreRec, createMode=createMode, editMode=editMode, shapeLineColor=shapeLineColor, shapeFillColor=shapeFillColor, zoom=zoom, zoomIn=zoomIn, zoomOut=zoomOut, zoomOrg=zoomOrg, fitWindow=fitWindow, fitWidth=fitWidth, 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), + 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, 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, createMode, editMode), + onLoadActive=(create, createpoly, createMode, editMode), onShapesPresent=(hideAll, showAll)) # menus @@ -574,7 +600,7 @@ class MainWindow(QMainWindow): self.autoSaveOption.triggered.connect(self.autoSaveFunc) addActions(self.menus.file, - (opendir, open_dataset_dir, None, saveLabel, saveRec, self.autoSaveOption, None, resetAll, deleteImg, + (opendir, open_dataset_dir, None, saveLabel, saveRec, exportJSON, self.autoSaveOption, None, resetAll, deleteImg, quit)) addActions(self.menus.help, (showKeys, showSteps, showInfo)) @@ -695,6 +721,7 @@ class MainWindow(QMainWindow): self.dirty = False self.actions.save.setEnabled(False) self.actions.create.setEnabled(True) + self.actions.createpoly.setEnabled(True) def toggleActions(self, value=True): """Enable/Disable widgets which depend on an opened image.""" @@ -780,6 +807,7 @@ class MainWindow(QMainWindow): assert self.beginner() self.canvas.setEditing(False) self.actions.create.setEnabled(False) + self.actions.createpoly.setEnabled(False) self.canvas.fourpoint = False def createPolygon(self): @@ -787,10 +815,10 @@ class MainWindow(QMainWindow): self.canvas.setEditing(False) self.canvas.fourpoint = True self.actions.create.setEnabled(False) + self.actions.createpoly.setEnabled(False) self.actions.undoLastPoint.setEnabled(True) def rotateImg(self, filename, k, _value): - self.actions.rotateRight.setEnabled(_value) pix = cv2.imread(filename) pix = np.rot90(pix, k) @@ -831,6 +859,7 @@ class MainWindow(QMainWindow): self.canvas.setEditing(True) self.canvas.restoreCursor() self.actions.create.setEnabled(True) + self.actions.createpoly.setEnabled(True) def toggleDrawMode(self, edit=True): self.canvas.setEditing(edit) @@ -992,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) @@ -1216,6 +1246,7 @@ class MainWindow(QMainWindow): if self.beginner(): # Switch to edit mode. self.canvas.setEditing(True) self.actions.create.setEnabled(True) + self.actions.createpoly.setEnabled(True) self.actions.undoLastPoint.setEnabled(False) self.actions.undo.setEnabled(True) else: @@ -1656,8 +1687,10 @@ class MainWindow(QMainWindow): self.haveAutoReced = False self.AutoRecognition.setEnabled(True) self.reRecogButton.setEnabled(True) + self.tableRecButton.setEnabled(True) self.actions.AutoRec.setEnabled(True) self.actions.reRec.setEnabled(True) + self.actions.tableRec.setEnabled(True) self.actions.open_dataset_dir.setEnabled(True) self.actions.rotateLeft.setEnabled(True) self.actions.rotateRight.setEnabled(True) @@ -1757,6 +1790,7 @@ class MainWindow(QMainWindow): self.openNextImg() self.actions.saveRec.setEnabled(True) self.actions.saveLabel.setEnabled(True) + self.actions.exportJSON.setEnabled(True) elif mode == 'Auto': if annotationFilePath and self.saveLabels(annotationFilePath, mode=mode): @@ -2083,6 +2117,266 @@ class MainWindow(QMainWindow): self.singleLabel(shape) self.setDirty() + def TableRecognition(self): + ''' + Table Recegnition + ''' + from paddleocr.ppstructure.table.predict_table import to_excel + + import time + + start = time.time() + img = cv2.imread(self.filePath) + res = self.table_ocr(img, return_ocr_result_in_table=True) + + TableRec_excel_dir = self.lastOpenDir + '/tableRec_excel_output/' + os.makedirs(TableRec_excel_dir, exist_ok=True) + filename = os.path.basename(self.filePath) + excel_path = TableRec_excel_dir + '{}.xlsx'.format(filename) + + if res is None: + msg = 'Can not recognise the table in ' + self.filePath + '. Please change manually' + QMessageBox.information(self, "Information", msg) + to_excel('', excel_path) # create an empty excel + return + + # save res + # ONLY SUPPORT ONE TABLE in one image + hasTable = False + for region in res: + if region['type'] == 'Table': + if region['res']['boxes'] is None: + msg = 'Can not recognise the detection box in ' + self.filePath + '. Please change manually' + QMessageBox.information(self, "Information", msg) + to_excel('', excel_path) # create an empty excel + return + hasTable = True + # save table ocr result on PPOCRLabel + # clear all old annotaions before saving result + self.itemsToShapes.clear() + self.shapesToItems.clear() + self.itemsToShapesbox.clear() # ADD + self.shapesToItemsbox.clear() + self.labelList.clear() + self.BoxList.clear() + self.result_dic = [] + self.result_dic_locked = [] + + shapes = [] + result_len = len(region['res']['boxes']) + for i in range(result_len): + bbox = np.array(region['res']['boxes'][i]) + rec_text = region['res']['rec_res'][i][0] + + # polys to rectangles + x1, y1 = np.min(bbox[:, 0]), np.min(bbox[:, 1]) + x2, y2 = np.max(bbox[:, 0]), np.max(bbox[:, 1]) + rext_bbox = [[x1, y1], [x2, y1], [x2, y2], [x1, y2]] + + # save bbox to shape + shape = Shape(label=rec_text, line_color=DEFAULT_LINE_COLOR, key_cls=None) + for point in rext_bbox: + x, y = point + # Ensure the labels are within the bounds of the image. + # If not, fix them. + x, y, snapped = self.canvas.snapPointToCanvas(x, y) + shape.addPoint(QPointF(x, y)) + shape.difficult = False + # shape.locked = False + shape.close() + self.addLabel(shape) + shapes.append(shape) + self.setDirty() + self.canvas.loadShapes(shapes) + + # save HTML result to excel + try: + to_excel(region['res']['html'], excel_path) + except: + print('Can not save excel file, maybe Permission denied (.xlsx is being occupied)') + break + + if not hasTable: + msg = 'Can not recognise the table in ' + self.filePath + '. Please change manually' + QMessageBox.information(self, "Information", msg) + to_excel('', excel_path) # create an empty excel + return + + # automatically open excel annotation file + try: + import win32com.client + except: + print("CANNOT OPEN .xlsx. It could be one of the following reasons: " \ + "Only support Windows | No python win32com") + + try: + xl = win32com.client.Dispatch("Excel.Application") + xl.Visible = True + xl.Workbooks.Open(excel_path) + except: + print("CANNOT OPEN .xlsx. It could be the following reasons: " \ + ".xlsx is not existed") + + print('time cost: ', time.time() - start) + + def cellreRecognition(self): + ''' + 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 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 + + # 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] != '': + 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] + else: + 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 libs.dataPartitionDialog import DataPartitionDialog + + if self.lang == 'ch': + QMessageBox.information(self, "Information", "导出JSON前请保存所有图像的标注且关闭EXCEL!!!!!!!!!!!!") + else: + QMessageBox.information(self, "Information", "Please save all the annotations and close the EXCEL before exporting JSON!!!!!!!!!!!!") + + # automatically save annotations + self.saveLabelFile() + + # load box annotations + labeldict = {} + if not os.path.exists(self.PPlabelpath): + msg = 'ERROR, Can not find Label.txt' + QMessageBox.information(self, "Information", msg) + return + else: + with open(self.PPlabelpath, '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] = [] + + # if len(labeldict) != len(csv_paths): + # msg = 'ERROR, box label and excel label are not in the same number\n' + \ + # 'box label: ' + str(len(labeldict)) + '\n' + \ + # 'excel label: ' + str(len(csv_paths)) + '\n' + \ + # 'Please check the label.txt and tableRec_excel_output\n' + # QMessageBox.information(self, "Information", msg) + # return + + # data partition user input + partitionDialog = DataPartitionDialog() + partitionDialog.exec() + if partitionDialog.getStatus() == False: + return + + train_split, val_split, test_split = partitionDialog.getDataPartition() + # check validate + if train_split + val_split + test_split > 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)) + print('Data partition: train:', train_id, + 'validation:', val_id - train_id, + 'test:', len(labeldict) - val_id) + + TableRec_excel_dir = os.path.join(self.lastOpenDir, 'tableRec_excel_output') + json_results = [] + imgid = 0 + for image_path in labeldict.keys(): + # load csv annotations + filename = os.path.basename(image_path) + csv_path = os.path.join(TableRec_excel_dir, filename + '.xlsx') + if not os.path.exists(csv_path): + msg = 'ERROR, Can not find ' + csv_path + QMessageBox.information(self, "Information", msg) + return + + # read xlsx file, convert to HTML + xd = pd.ExcelFile(csv_path) + df = xd.parse() + structure = df.to_html() + + # load box annotations + cells = [] + for anno in labeldict[image_path]: + tokens = list(anno['transcription']) + obb = anno['points'] + hbb = OBB2HBB(np.array(obb)).tolist() + cells.append({'tokens': tokens, 'bbox': hbb}) + + # data split + if imgid < train_id: + split = 'train' + elif imgid < val_id: + split = 'val' + else: + split = 'test' + + # save dict + html = {'structure': {'tokens': structure}, 'cell': cells} + json_results.append({'filename': filename, 'split': split, 'imgid': imgid, 'html': html}) + imgid += 1 + + # save json + 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) + QMessageBox.information(self, "Information", msg) + def autolcm(self): vbox = QVBoxLayout() hbox = QHBoxLayout() @@ -2122,6 +2416,12 @@ class MainWindow(QMainWindow): 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()]) + del self.table_ocr + self.table_ocr = PPStructure(use_pdserving=False, + use_gpu=False, + lang=lg_idx[self.comboBox.currentText()], + layout=False, + show_log=False) self.dialog.close() def cancel(self): @@ -2140,6 +2440,7 @@ class MainWindow(QMainWindow): self.fileStatedict[file] = 1 self.actions.saveLabel.setEnabled(True) self.actions.saveRec.setEnabled(True) + self.actions.exportJSON.setEnabled(True) def saveFilestate(self): with open(self.fileStatepath, 'w', encoding='utf-8') as f: diff --git a/PPOCRLabel/libs/dataPartitionDialog.py b/PPOCRLabel/libs/dataPartitionDialog.py new file mode 100644 index 0000000000..051b9722d5 --- /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/libs/utils.py b/PPOCRLabel/libs/utils.py index 2510520caa..c49b506882 100644 --- a/PPOCRLabel/libs/utils.py +++ b/PPOCRLabel/libs/utils.py @@ -161,6 +161,33 @@ def get_rotate_crop_image(img, points): print(e) +def boxPad(box, imgShape, pad : int) -> np.array: + """ + Pad a box with [pad] pixels on each side. + """ + box = np.array(box, dtype=np.int32) + box[0][0], box[0][1] = box[0][0] - pad, box[0][1] - pad + box[1][0], box[1][1] = box[1][0] + pad, box[1][1] - pad + box[2][0], box[2][1] = box[2][0] + pad, box[2][1] + pad + box[3][0], box[3][1] = box[3][0] - pad, box[3][1] + pad + h, w, _ = imgShape + box[:,0] = np.clip(box[:,0], 0, w) + box[:,1] = np.clip(box[:,1], 0, h) + return box + + +def OBB2HBB(obb) -> np.array: + """ + Convert Oriented Bounding Box to Horizontal Bounding Box. + """ + hbb = np.zeros(4, dtype=np.int32) + hbb[0] = min(obb[:, 0]) + hbb[1] = min(obb[:, 1]) + hbb[2] = max(obb[:, 0]) + hbb[3] = max(obb[:, 1]) + return hbb + + def stepsInfo(lang='en'): if lang == 'ch': msg = "1. 安装与运行：使用上述命令安装与运行程序。\n" \ diff --git a/PPOCRLabel/resources/strings/strings-en.properties b/PPOCRLabel/resources/strings/strings-en.properties index 3c4eda65a3..7ba9af4c33 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 @@ -110,3 +110,6 @@ lockBoxDetail=Lock selected box/Unlock all box keyListTitle=Key List keyDialogTip=Enter object label keyChange=Change Box Key +TableRecognition=Table Recognition +cellreRecognition=Cell Re-Recognition +exportJSON=export JSON(PubTabNet) diff --git a/PPOCRLabel/resources/strings/strings-zh-CN.properties b/PPOCRLabel/resources/strings/strings-zh-CN.properties index a7c30368b8..308974ef29 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度 @@ -109,4 +109,7 @@ lockBox=锁定框/解除锁定框 lockBoxDetail=若当前没有框处于锁定状态则锁定选中的框，若存在锁定框则解除所有锁定框的锁定状态 keyListTitle=关键词列表 keyDialogTip=请输入类型名称 -keyChange=更改Box关键字类别 \ No newline at end of file +keyChange=更改Box关键字类别 +TableRecognition=表格识别 +cellreRecognition=单元格重识别 +exportJSON=导出表格JSON标注 \ No newline at end of file From 15be3c54d80a76aca4518b9f800378b4093a1c8b Mon Sep 17 00:00:00 2001 From: Leif <4603009@qq.com> Date: Thu, 5 May 2022 11:26:55 +0800 Subject: [PATCH 3/7] Auto stash before merge of "table" and "origin/new" --- PPOCRLabel/PPOCRLabel.py | 39 +++++++++++++++++++++++---------------- 1 file changed, 23 insertions(+), 16 deletions(-) diff --git a/PPOCRLabel/PPOCRLabel.py b/PPOCRLabel/PPOCRLabel.py index 4d6c0d7770..6729c57487 100644 --- a/PPOCRLabel/PPOCRLabel.py +++ b/PPOCRLabel/PPOCRLabel.py @@ -2131,7 +2131,7 @@ class MainWindow(QMainWindow): TableRec_excel_dir = self.lastOpenDir + '/tableRec_excel_output/' os.makedirs(TableRec_excel_dir, exist_ok=True) - filename = os.path.basename(self.filePath) + filename, _ = os.path.splitext(os.path.basename(self.filePath)) excel_path = TableRec_excel_dir + '{}.xlsx'.format(filename) if res is None: @@ -2203,19 +2203,26 @@ class MainWindow(QMainWindow): return # automatically open excel annotation file - try: - import win32com.client - except: - print("CANNOT OPEN .xlsx. It could be one of the following reasons: " \ - "Only support Windows | No python win32com") + if platform.system() == 'Windows': + try: + import win32com.client + except: + print("CANNOT OPEN .xlsx. It could be one of the following reasons: " \ + "Only support Windows | No python win32com") - try: - xl = win32com.client.Dispatch("Excel.Application") - xl.Visible = True - xl.Workbooks.Open(excel_path) - except: - print("CANNOT OPEN .xlsx. It could be the following reasons: " \ - ".xlsx is not existed") + try: + xl = win32com.client.Dispatch("Excel.Application") + xl.Visible = True + xl.Workbooks.Open(excel_path) + # excelEx = "You need to show the excel executable at this point" + # subprocess.Popen([excelEx, excel_path]) + + # os.startfile(excel_path) + except: + print("CANNOT OPEN .xlsx. It could be the following reasons: " \ + ".xlsx is not existed") + else: + os.system('open ' + os.path.normpath(excel_path)) print('time cost: ', time.time() - start) @@ -2337,7 +2344,7 @@ class MainWindow(QMainWindow): imgid = 0 for image_path in labeldict.keys(): # load csv annotations - filename = os.path.basename(image_path) + filename, _ = os.path.splitext(os.path.basename(image_path)) csv_path = os.path.join(TableRec_excel_dir, filename + '.xlsx') if not os.path.exists(csv_path): msg = 'ERROR, Can not find ' + csv_path @@ -2347,7 +2354,7 @@ class MainWindow(QMainWindow): # read xlsx file, convert to HTML xd = pd.ExcelFile(csv_path) df = xd.parse() - structure = df.to_html() + structure = df.to_html(index = False) # load box annotations cells = [] @@ -2371,7 +2378,7 @@ class MainWindow(QMainWindow): imgid += 1 # save json - with open("{}/annotation.json".format(self.lastOpenDir), "w") as fid: + with open("{}/annotation.json".format(self.lastOpenDir), "w", encoding='utf-8') as fid: fid.write(json.dumps(json_results)) msg = 'JSON sucessfully saved in {}/annotation.json'.format(self.lastOpenDir) From 6413fb1f83a47aa0125b6d8fc235c0c213585dfb Mon Sep 17 00:00:00 2001 From: whj_dark Date: Wed, 4 May 2022 13:59:44 +0800 Subject: [PATCH 4/7] new --- PPOCRLabel/PPOCRLabel.py | 328 +++++++++++++++++- PPOCRLabel/libs/dataPartitionDialog.py | 113 ++++++ PPOCRLabel/libs/utils.py | 27 ++ .../resources/strings/strings-en.properties | 5 +- .../strings/strings-zh-CN.properties | 7 +- 5 files changed, 462 insertions(+), 18 deletions(-) create mode 100644 PPOCRLabel/libs/dataPartitionDialog.py diff --git a/PPOCRLabel/PPOCRLabel.py b/PPOCRLabel/PPOCRLabel.py index 523661b635..703c4e40cd 100644 --- a/PPOCRLabel/PPOCRLabel.py +++ b/PPOCRLabel/PPOCRLabel.py @@ -26,7 +26,7 @@ from functools import partial 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, \ + QSlider, QGraphicsOpacityEffect, QMessageBox, QListView, QScrollArea, QWidgetAction, QApplication, QLabel, QGridLayout, \ QFileDialog, QListWidgetItem, QComboBox, QDialog __dir__ = os.path.dirname(os.path.abspath(__file__)) @@ -36,7 +36,7 @@ sys.path.append(os.path.abspath(os.path.join(__dir__, '../..'))) sys.path.append(os.path.abspath(os.path.join(__dir__, '../PaddleOCR'))) sys.path.append("..") -from paddleocr import PaddleOCR +from paddleocr import PaddleOCR, PPStructure from libs.constants import * from libs.utils import * from libs.labelColor import label_colormap @@ -100,9 +100,15 @@ class MainWindow(QMainWindow): use_gpu=gpu, lang=lang, show_log=False) + self.table_ocr = PPStructure(use_pdserving=False, + use_gpu=gpu, + lang=lang, + layout=False, + show_log=False) if os.path.exists('./data/paddle.png'): result = self.ocr.ocr('./data/paddle.png', cls=True, det=True) + result = self.table_ocr('./data/paddle.png', return_ocr_result_in_table=True) # For loading all image under a directory self.mImgList = [] @@ -196,16 +202,25 @@ class MainWindow(QMainWindow): self.reRecogButton.setIcon(newIcon('reRec', 30)) self.reRecogButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon) + self.tableRecButton = QToolButton() + self.tableRecButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon) + self.newButton = QToolButton() self.newButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon) + self.createpolyButton = QToolButton() + self.createpolyButton.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) + leftTopToolBox = QGridLayout() + leftTopToolBox.addWidget(self.newButton, 0, 0, 1, 1) + leftTopToolBox.addWidget(self.createpolyButton, 0, 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) @@ -446,13 +461,22 @@ class MainWindow(QMainWindow): 'Ctrl+R', 'reRec', getStr('singleRe'), enabled=False) createpoly = action(getStr('creatPolygon'), self.createPolygon, - 'q', 'new', getStr('creatPolygon'), enabled=True) + 'q', 'new', getStr('creatPolygon'), enabled=False) + + tableRec = action(getStr('TableRecognition'), self.TableRecognition, + '', 'Auto', getStr('TableRecognition'), enabled=False) + + cellreRec = action(getStr('cellreRecognition'), self.cellreRecognition, + '', 'reRec', getStr('cellreRecognition'), enabled=False) saveRec = action(getStr('saveRec'), self.saveRecResult, '', 'save', getStr('saveRec'), enabled=False) saveLabel = action(getStr('saveLabel'), self.saveLabelFile, # 'Ctrl+S', 'save', getStr('saveLabel'), enabled=False) + + exportJSON = action(getStr('exportJSON'), self.exportJSON, + '', 'save', getStr('exportJSON'), enabled=False) undoLastPoint = action(getStr("undoLastPoint"), self.canvas.undoLastPoint, 'Ctrl+Z', "undo", getStr("undoLastPoint"), enabled=False) @@ -474,10 +498,12 @@ class MainWindow(QMainWindow): self.editButton.setDefaultAction(edit) self.newButton.setDefaultAction(create) + self.createpolyButton.setDefaultAction(createpoly) self.DelButton.setDefaultAction(deleteImg) self.SaveButton.setDefaultAction(save) self.AutoRecognition.setDefaultAction(AutoRec) self.reRecogButton.setDefaultAction(reRec) + self.tableRecButton.setDefaultAction(tableRec) # self.preButton.setDefaultAction(openPrevImg) # self.nextButton.setDefaultAction(openNextImg) @@ -523,25 +549,25 @@ class MainWindow(QMainWindow): # 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, + lineColor=color1, create=create, createpoly=createpoly, tableRec=tableRec, delete=delete, edit=edit, copy=copy, + saveRec=saveRec, singleRere=singleRere, AutoRec=AutoRec, reRec=reRec, cellreRec=cellreRec, createMode=createMode, editMode=editMode, shapeLineColor=shapeLineColor, shapeFillColor=shapeFillColor, zoom=zoom, zoomIn=zoomIn, zoomOut=zoomOut, zoomOrg=zoomOrg, fitWindow=fitWindow, fitWidth=fitWidth, 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), + 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, 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, createMode, editMode), + onLoadActive=(create, createpoly, createMode, editMode), onShapesPresent=(hideAll, showAll)) # menus @@ -574,7 +600,7 @@ class MainWindow(QMainWindow): self.autoSaveOption.triggered.connect(self.autoSaveFunc) addActions(self.menus.file, - (opendir, open_dataset_dir, None, saveLabel, saveRec, self.autoSaveOption, None, resetAll, deleteImg, + (opendir, open_dataset_dir, None, saveLabel, saveRec, exportJSON, self.autoSaveOption, None, resetAll, deleteImg, quit)) addActions(self.menus.help, (showKeys, showSteps, showInfo)) @@ -695,6 +721,7 @@ class MainWindow(QMainWindow): self.dirty = False self.actions.save.setEnabled(False) self.actions.create.setEnabled(True) + self.actions.createpoly.setEnabled(True) def toggleActions(self, value=True): """Enable/Disable widgets which depend on an opened image.""" @@ -780,6 +807,7 @@ class MainWindow(QMainWindow): assert self.beginner() self.canvas.setEditing(False) self.actions.create.setEnabled(False) + self.actions.createpoly.setEnabled(False) self.canvas.fourpoint = False def createPolygon(self): @@ -787,10 +815,10 @@ class MainWindow(QMainWindow): self.canvas.setEditing(False) self.canvas.fourpoint = True self.actions.create.setEnabled(False) + self.actions.createpoly.setEnabled(False) self.actions.undoLastPoint.setEnabled(True) def rotateImg(self, filename, k, _value): - self.actions.rotateRight.setEnabled(_value) pix = cv2.imread(filename) pix = np.rot90(pix, k) @@ -831,6 +859,7 @@ class MainWindow(QMainWindow): self.canvas.setEditing(True) self.canvas.restoreCursor() self.actions.create.setEnabled(True) + self.actions.createpoly.setEnabled(True) def toggleDrawMode(self, edit=True): self.canvas.setEditing(edit) @@ -992,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) @@ -1216,6 +1246,7 @@ class MainWindow(QMainWindow): if self.beginner(): # Switch to edit mode. self.canvas.setEditing(True) self.actions.create.setEnabled(True) + self.actions.createpoly.setEnabled(True) self.actions.undoLastPoint.setEnabled(False) self.actions.undo.setEnabled(True) else: @@ -1656,8 +1687,10 @@ class MainWindow(QMainWindow): self.haveAutoReced = False self.AutoRecognition.setEnabled(True) self.reRecogButton.setEnabled(True) + self.tableRecButton.setEnabled(True) self.actions.AutoRec.setEnabled(True) self.actions.reRec.setEnabled(True) + self.actions.tableRec.setEnabled(True) self.actions.open_dataset_dir.setEnabled(True) self.actions.rotateLeft.setEnabled(True) self.actions.rotateRight.setEnabled(True) @@ -1757,6 +1790,7 @@ class MainWindow(QMainWindow): self.openNextImg() self.actions.saveRec.setEnabled(True) self.actions.saveLabel.setEnabled(True) + self.actions.exportJSON.setEnabled(True) elif mode == 'Auto': if annotationFilePath and self.saveLabels(annotationFilePath, mode=mode): @@ -2083,6 +2117,263 @@ class MainWindow(QMainWindow): self.singleLabel(shape) self.setDirty() + def TableRecognition(self): + ''' + Table Recegnition + ''' + from paddleocr.ppstructure.table.predict_table import to_excel + + import time + + start = time.time() + img = cv2.imread(self.filePath) + res = self.table_ocr(img, return_ocr_result_in_table=True) + + TableRec_excel_dir = self.lastOpenDir + '/tableRec_excel_output/' + os.makedirs(TableRec_excel_dir, exist_ok=True) + filename = os.path.basename(self.filePath) + excel_path = TableRec_excel_dir + '{}.xlsx'.format(filename) + + if res is None: + msg = 'Can not recognise the table in ' + self.filePath + '. Please change manually' + QMessageBox.information(self, "Information", msg) + to_excel('', excel_path) # create an empty excel + return + + # save res + # ONLY SUPPORT ONE TABLE in one image + hasTable = False + for region in res: + if region['type'] == 'Table': + if region['res']['boxes'] is None: + msg = 'Can not recognise the detection box in ' + self.filePath + '. Please change manually' + QMessageBox.information(self, "Information", msg) + to_excel('', excel_path) # create an empty excel + return + hasTable = True + # save table ocr result on PPOCRLabel + # clear all old annotaions before saving result + self.itemsToShapes.clear() + self.shapesToItems.clear() + self.itemsToShapesbox.clear() # ADD + self.shapesToItemsbox.clear() + self.labelList.clear() + self.BoxList.clear() + self.result_dic = [] + self.result_dic_locked = [] + + shapes = [] + result_len = len(region['res']['boxes']) + for i in range(result_len): + bbox = np.array(region['res']['boxes'][i]) + rec_text = region['res']['rec_res'][i][0] + + # polys to rectangles + x1, y1 = np.min(bbox[:, 0]), np.min(bbox[:, 1]) + x2, y2 = np.max(bbox[:, 0]), np.max(bbox[:, 1]) + rext_bbox = [[x1, y1], [x2, y1], [x2, y2], [x1, y2]] + + # save bbox to shape + shape = Shape(label=rec_text, line_color=DEFAULT_LINE_COLOR, key_cls=None) + for point in rext_bbox: + x, y = point + # Ensure the labels are within the bounds of the image. + # If not, fix them. + x, y, snapped = self.canvas.snapPointToCanvas(x, y) + shape.addPoint(QPointF(x, y)) + shape.difficult = False + # shape.locked = False + shape.close() + self.addLabel(shape) + shapes.append(shape) + self.setDirty() + self.canvas.loadShapes(shapes) + + # save HTML result to excel + try: + to_excel(region['res']['html'], excel_path) + except: + print('Can not save excel file, maybe Permission denied (.xlsx is being occupied)') + break + + if not hasTable: + msg = 'Can not recognise the table in ' + self.filePath + '. Please change manually' + QMessageBox.information(self, "Information", msg) + to_excel('', excel_path) # create an empty excel + return + + # automatically open excel annotation file + try: + import win32com.client + except: + print("CANNOT OPEN .xlsx. It could be one of the following reasons: " \ + "Only support Windows | No python win32com") + + try: + xl = win32com.client.Dispatch("Excel.Application") + xl.Visible = True + xl.Workbooks.Open(excel_path) + except: + print("CANNOT OPEN .xlsx. It could be the following reasons: " \ + ".xlsx is not existed") + + print('time cost: ', time.time() - start) + + def cellreRecognition(self): + ''' + 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 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 + + # 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] != '': + 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] + else: + 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 libs.dataPartitionDialog import DataPartitionDialog + + # data partition user input + partitionDialog = DataPartitionDialog(parent=self) + partitionDialog.exec() + if partitionDialog.getStatus() == False: + return + + # automatically save annotations + self.saveFilestate() + self.savePPlabel(mode='auto') + + # load box annotations + labeldict = {} + if not os.path.exists(self.PPlabelpath): + msg = 'ERROR, Can not find Label.txt' + QMessageBox.information(self, "Information", msg) + return + else: + with open(self.PPlabelpath, '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] = [] + + # if len(labeldict) != len(csv_paths): + # msg = 'ERROR, box label and excel label are not in the same number\n' + \ + # 'box label: ' + str(len(labeldict)) + '\n' + \ + # 'excel label: ' + str(len(csv_paths)) + '\n' + \ + # 'Please check the label.txt and tableRec_excel_output\n' + # QMessageBox.information(self, "Information", msg) + # return + + + train_split, val_split, test_split = partitionDialog.getDataPartition() + # check validate + if train_split + val_split + test_split > 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)) + print('Data partition: train:', train_id, + 'validation:', val_id - train_id, + 'test:', len(labeldict) - val_id) + + TableRec_excel_dir = os.path.join(self.lastOpenDir, 'tableRec_excel_output') + json_results = [] + imgid = 0 + for image_path in labeldict.keys(): + # load csv annotations + filename = os.path.basename(image_path) + csv_path = os.path.join(TableRec_excel_dir, filename + '.xlsx') + if not os.path.exists(csv_path): + msg = 'ERROR, Can not find ' + csv_path + QMessageBox.information(self, "Information", msg) + return + + # read xlsx file, convert to HTML + xd = pd.ExcelFile(csv_path) + df = xd.parse() + structure = df.to_html() + + # load box annotations + cells = [] + for anno in labeldict[image_path]: + tokens = list(anno['transcription']) + obb = anno['points'] + hbb = OBB2HBB(np.array(obb)).tolist() + cells.append({'tokens': tokens, 'bbox': hbb}) + + # data split + if imgid < train_id: + split = 'train' + elif imgid < val_id: + split = 'val' + else: + split = 'test' + + # save dict + html = {'structure': {'tokens': structure}, 'cell': cells} + json_results.append({'filename': filename, 'split': split, 'imgid': imgid, 'html': html}) + imgid += 1 + + # save json + with open("{}/annotation.json".format(self.lastOpenDir), "w", encoding='utf-8') as fid: + fid.write(json.dumps(json_results, ensure_ascii=False)) + + msg = 'JSON sucessfully saved in {}/annotation.json'.format(self.lastOpenDir) + QMessageBox.information(self, "Information", msg) + def autolcm(self): vbox = QVBoxLayout() hbox = QHBoxLayout() @@ -2122,6 +2413,12 @@ class MainWindow(QMainWindow): 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()]) + del self.table_ocr + self.table_ocr = PPStructure(use_pdserving=False, + use_gpu=False, + lang=lg_idx[self.comboBox.currentText()], + layout=False, + show_log=False) self.dialog.close() def cancel(self): @@ -2140,6 +2437,7 @@ class MainWindow(QMainWindow): self.fileStatedict[file] = 1 self.actions.saveLabel.setEnabled(True) self.actions.saveRec.setEnabled(True) + self.actions.exportJSON.setEnabled(True) def saveFilestate(self): with open(self.fileStatepath, 'w', encoding='utf-8') as f: diff --git a/PPOCRLabel/libs/dataPartitionDialog.py b/PPOCRLabel/libs/dataPartitionDialog.py new file mode 100644 index 0000000000..33bd491552 --- /dev/null +++ b/PPOCRLabel/libs/dataPartitionDialog.py @@ -0,0 +1,113 @@ +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, parent=None): + super().__init__() + self.parnet = parent + 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 + + if self.parnet.lang == 'ch': + msg = "导出JSON前请保存所有图像的标注且关闭EXCEL!" + else: + msg = "Please save all the annotations and close the EXCEL before exporting JSON!" + + info_msg = QLabel(msg, self) + info_msg.setWordWrap(True) + info_msg.setStyleSheet("color: red") + info_msg.setFont(QFont('Arial', 12)) + + 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(info_msg, 0, 0, 1, 2) + gridlayout.addWidget(train_lbl, 1, 0) + gridlayout.addWidget(val_lbl, 2, 0) + gridlayout.addWidget(test_lbl, 3, 0) + gridlayout.addWidget(self.train_input, 1, 1) + gridlayout.addWidget(self.val_input, 2, 1) + gridlayout.addWidget(self.test_input, 3, 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, 4, 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/libs/utils.py b/PPOCRLabel/libs/utils.py index 2510520caa..c49b506882 100644 --- a/PPOCRLabel/libs/utils.py +++ b/PPOCRLabel/libs/utils.py @@ -161,6 +161,33 @@ def get_rotate_crop_image(img, points): print(e) +def boxPad(box, imgShape, pad : int) -> np.array: + """ + Pad a box with [pad] pixels on each side. + """ + box = np.array(box, dtype=np.int32) + box[0][0], box[0][1] = box[0][0] - pad, box[0][1] - pad + box[1][0], box[1][1] = box[1][0] + pad, box[1][1] - pad + box[2][0], box[2][1] = box[2][0] + pad, box[2][1] + pad + box[3][0], box[3][1] = box[3][0] - pad, box[3][1] + pad + h, w, _ = imgShape + box[:,0] = np.clip(box[:,0], 0, w) + box[:,1] = np.clip(box[:,1], 0, h) + return box + + +def OBB2HBB(obb) -> np.array: + """ + Convert Oriented Bounding Box to Horizontal Bounding Box. + """ + hbb = np.zeros(4, dtype=np.int32) + hbb[0] = min(obb[:, 0]) + hbb[1] = min(obb[:, 1]) + hbb[2] = max(obb[:, 0]) + hbb[3] = max(obb[:, 1]) + return hbb + + def stepsInfo(lang='en'): if lang == 'ch': msg = "1. 安装与运行：使用上述命令安装与运行程序。\n" \ diff --git a/PPOCRLabel/resources/strings/strings-en.properties b/PPOCRLabel/resources/strings/strings-en.properties index 3c4eda65a3..7ba9af4c33 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 @@ -110,3 +110,6 @@ lockBoxDetail=Lock selected box/Unlock all box keyListTitle=Key List keyDialogTip=Enter object label keyChange=Change Box Key +TableRecognition=Table Recognition +cellreRecognition=Cell Re-Recognition +exportJSON=export JSON(PubTabNet) diff --git a/PPOCRLabel/resources/strings/strings-zh-CN.properties b/PPOCRLabel/resources/strings/strings-zh-CN.properties index a7c30368b8..308974ef29 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度 @@ -109,4 +109,7 @@ lockBox=锁定框/解除锁定框 lockBoxDetail=若当前没有框处于锁定状态则锁定选中的框，若存在锁定框则解除所有锁定框的锁定状态 keyListTitle=关键词列表 keyDialogTip=请输入类型名称 -keyChange=更改Box关键字类别 \ No newline at end of file +keyChange=更改Box关键字类别 +TableRecognition=表格识别 +cellreRecognition=单元格重识别 +exportJSON=导出表格JSON标注 \ No newline at end of file From 1e6af3bb886107639de901c1cecd7061f18480a8 Mon Sep 17 00:00:00 2001 From: Leif <4603009@qq.com> Date: Thu, 5 May 2022 22:36:59 +0800 Subject: [PATCH 5/7] merge from new merge from new --- PPOCRLabel/PPOCRLabel.py | 21 +++++++++---------- PPOCRLabel/libs/dataPartitionDialog.py | 28 ++++++++++++++++++-------- 2 files changed, 29 insertions(+), 20 deletions(-) diff --git a/PPOCRLabel/PPOCRLabel.py b/PPOCRLabel/PPOCRLabel.py index 6729c57487..534e0cd882 100644 --- a/PPOCRLabel/PPOCRLabel.py +++ b/PPOCRLabel/PPOCRLabel.py @@ -2285,13 +2285,15 @@ class MainWindow(QMainWindow): import pandas as pd from libs.dataPartitionDialog import DataPartitionDialog - if self.lang == 'ch': - QMessageBox.information(self, "Information", "导出JSON前请保存所有图像的标注且关闭EXCEL!!!!!!!!!!!!") - else: - QMessageBox.information(self, "Information", "Please save all the annotations and close the EXCEL before exporting JSON!!!!!!!!!!!!") + # data partition user input + partitionDialog = DataPartitionDialog(parent=self) + partitionDialog.exec() + if partitionDialog.getStatus() == False: + return # automatically save annotations - self.saveLabelFile() + self.saveFilestate() + self.savePPlabel(mode='auto') # load box annotations labeldict = {} @@ -2318,12 +2320,7 @@ class MainWindow(QMainWindow): # 'Please check the label.txt and tableRec_excel_output\n' # QMessageBox.information(self, "Information", msg) # return - - # data partition user input - partitionDialog = DataPartitionDialog() - partitionDialog.exec() - if partitionDialog.getStatus() == False: - return + train_split, val_split, test_split = partitionDialog.getDataPartition() # check validate @@ -2379,7 +2376,7 @@ class MainWindow(QMainWindow): # save json with open("{}/annotation.json".format(self.lastOpenDir), "w", encoding='utf-8') as fid: - fid.write(json.dumps(json_results)) + fid.write(json.dumps(json_results, ensure_ascii=False)) msg = 'JSON sucessfully saved in {}/annotation.json'.format(self.lastOpenDir) QMessageBox.information(self, "Information", msg) diff --git a/PPOCRLabel/libs/dataPartitionDialog.py b/PPOCRLabel/libs/dataPartitionDialog.py index 051b9722d5..33bd491552 100644 --- a/PPOCRLabel/libs/dataPartitionDialog.py +++ b/PPOCRLabel/libs/dataPartitionDialog.py @@ -18,8 +18,9 @@ import numpy as np BB = QDialogButtonBox class DataPartitionDialog(QDialog): - def __init__(self): + def __init__(self, parent=None): super().__init__() + self.parnet = parent self.title = 'DATA PARTITION' self.train_ratio = 70 @@ -34,6 +35,16 @@ class DataPartitionDialog(QDialog): self.flag_accept = True + if self.parnet.lang == 'ch': + msg = "导出JSON前请保存所有图像的标注且关闭EXCEL!" + else: + msg = "Please save all the annotations and close the EXCEL before exporting JSON!" + + info_msg = QLabel(msg, self) + info_msg.setWordWrap(True) + info_msg.setStyleSheet("color: red") + info_msg.setFont(QFont('Arial', 12)) + train_lbl = QLabel('Train split: ', self) train_lbl.setFont(QFont('Arial', 15)) val_lbl = QLabel('Valid split: ', self) @@ -58,19 +69,20 @@ class DataPartitionDialog(QDialog): 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) + gridlayout.addWidget(info_msg, 0, 0, 1, 2) + gridlayout.addWidget(train_lbl, 1, 0) + gridlayout.addWidget(val_lbl, 2, 0) + gridlayout.addWidget(test_lbl, 3, 0) + gridlayout.addWidget(self.train_input, 1, 1) + gridlayout.addWidget(self.val_input, 2, 1) + gridlayout.addWidget(self.test_input, 3, 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) + gridlayout.addWidget(bb, 4, 0, 1, 2) self.setLayout(gridlayout) From f1d0c8a2ff1a5790f78749d7cd489d8d8337b463 Mon Sep 17 00:00:00 2001 From: Leif <4603009@qq.com> Date: Thu, 5 May 2022 22:37:11 +0800 Subject: [PATCH 6/7] Auto stash before merge of "table" and "origin/new" --- PPOCRLabel/PPOCRLabel.py | 19 +++++++++++++---- PPOCRLabel/libs/utils.py | 44 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 59 insertions(+), 4 deletions(-) diff --git a/PPOCRLabel/PPOCRLabel.py b/PPOCRLabel/PPOCRLabel.py index 534e0cd882..f80cf8fd0e 100644 --- a/PPOCRLabel/PPOCRLabel.py +++ b/PPOCRLabel/PPOCRLabel.py @@ -21,6 +21,7 @@ import os.path import platform import subprocess import sys +import xlrd from functools import partial from PyQt5.QtCore import QSize, Qt, QPoint, QByteArray, QTimer, QFileInfo, QPointF, QProcess @@ -2349,9 +2350,19 @@ class MainWindow(QMainWindow): return # read xlsx file, convert to HTML - xd = pd.ExcelFile(csv_path) - df = xd.parse() - structure = df.to_html(index = False) + # xd = pd.ExcelFile(csv_path) + # df = xd.parse() + # structure = df.to_html(index = False) + excel = xlrd.open_workbook(csv_path) + sheet0 = excel.sheet_by_index(0) # only sheet 0 + merged_cells = sheet0.merged_cells # (0,1,1,3) start row, end row, start col, end col + + html_list = [['td'] * sheet0.ncols for i in range(sheet0.nrows)] + + for merged in merged_cells: + html_list = expand_list(merged, html_list) + + token_list = convert_token(html_list) # load box annotations cells = [] @@ -2370,7 +2381,7 @@ class MainWindow(QMainWindow): split = 'test' # save dict - html = {'structure': {'tokens': structure}, 'cell': cells} + html = {'structure': {'tokens': token_list}, 'cell': cells} json_results.append({'filename': filename, 'split': split, 'imgid': imgid, 'html': html}) imgid += 1 diff --git a/PPOCRLabel/libs/utils.py b/PPOCRLabel/libs/utils.py index c49b506882..86a0336d45 100644 --- a/PPOCRLabel/libs/utils.py +++ b/PPOCRLabel/libs/utils.py @@ -188,6 +188,50 @@ def OBB2HBB(obb) -> np.array: return hbb +def expand_list(merged, html_list): + ''' + Fill blanks according to merged cells + ''' + sr, er, sc, ec = merged + for i in range(sr, er): + for j in range(sc, ec): + html_list[i][j] = None + html_list[sr][sc] = '' + if ec - sc > 1: + html_list[sr][sc] += " colspan={}".format(ec - sc) + if er - sr > 1: + html_list[sr][sc] += " rowspan={}".format(er - sr) + return html_list + + +def convert_token(html_list): + ''' + Convert raw html to label format + ''' + token_list = [""] + # final html list: + for row in html_list: + token_list.append("") + for col in row: + if col == None: + continue + elif col == 'td': + token_list.extend(["", ""]) + else: + token_list.append("") + token_list.append("") + token_list.append("") + + return token_list + + def stepsInfo(lang='en'): if lang == 'ch': msg = "1. 安装与运行：使用上述命令安装与运行程序。\n" \ From b657a3a684b82d0356e6dc51111fc688a32c40ea Mon Sep 17 00:00:00 2001 From: Leif <4603009@qq.com> Date: Fri, 6 May 2022 11:49:06 +0800 Subject: [PATCH 7/7] change label filename change label filename --- PPOCRLabel/PPOCRLabel.py | 2 +- PPOCRLabel/libs/utils.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/PPOCRLabel/PPOCRLabel.py b/PPOCRLabel/PPOCRLabel.py index f80cf8fd0e..5839ec19dc 100644 --- a/PPOCRLabel/PPOCRLabel.py +++ b/PPOCRLabel/PPOCRLabel.py @@ -2382,7 +2382,7 @@ class MainWindow(QMainWindow): # save dict html = {'structure': {'tokens': token_list}, 'cell': cells} - json_results.append({'filename': filename, 'split': split, 'imgid': imgid, 'html': html}) + json_results.append({'filename': os.path.basename(image_path), 'split': split, 'imgid': imgid, 'html': html}) imgid += 1 # save json diff --git a/PPOCRLabel/libs/utils.py b/PPOCRLabel/libs/utils.py index 86a0336d45..bf54700488 100644 --- a/PPOCRLabel/libs/utils.py +++ b/PPOCRLabel/libs/utils.py @@ -225,7 +225,7 @@ def convert_token(html_list): if 'rowspan' in col: _, n = col.split('rowspan=') token_list.append(" rowspan=\"{}\"".format(n)) - token_list.append(">") + token_list.extend([">", ""]) token_list.append("") token_list.append("")