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 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__, '../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 @@ def __init__(self, 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 @@ def __init__(self, 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 @@ def __init__(self, '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 @@ def __init__(self, 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 @@ def __init__(self, # 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 @@ def __init__(self, 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 @@ def __init__(self, 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 @@ def setClean(self): 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 @@ def createShape(self): 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 @@ def createPolygon(self): 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 @@ def toggleDrawingSensitive(self, drawing=True): 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 @@ def newShape(self, value=True): 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 @@ def importDirImages(self, dirpath, isDelete=False): 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 @@ def _saveFile(self, annotationFilePath, mode='Manual'): 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 @@ def singleRerecognition(self): 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 @@ def modelChoose(self): 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 @@ def loadFilestate(self, saveDir): 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 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__, '../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 @@ def __init__(self, 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 @@ def __init__(self, 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 @@ def __init__(self, '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 @@ def __init__(self, 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 @@ def __init__(self, # 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 @@ def __init__(self, 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 @@ def setClean(self): 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 @@ def createShape(self): 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 @@ def createPolygon(self): 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 @@ def toggleDrawingSensitive(self, drawing=True): 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 @@ def shapeSelectionChanged(self, selected_shapes): 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 @@ def newShape(self, value=True): 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 @@ def importDirImages(self, dirpath, isDelete=False): 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 @@ def _saveFile(self, annotationFilePath, mode='Manual'): 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 @@ def singleRerecognition(self): 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 @@ def modelChoose(self): 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 @@ def loadFilestate(self, saveDir): 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 @@ def TableRecognition(self): 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 @@ def TableRecognition(self): 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 @@ def exportJSON(self): 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 @@ def exportJSON(self): # 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 @@ def exportJSON(self): 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 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__, '../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 @@ def __init__(self, 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 @@ def __init__(self, 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 @@ def __init__(self, '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 @@ def __init__(self, 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 @@ def __init__(self, # 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 @@ def __init__(self, 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 @@ def setClean(self): 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 @@ def createShape(self): 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 @@ def createPolygon(self): 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 @@ def toggleDrawingSensitive(self, drawing=True): 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 @@ def shapeSelectionChanged(self, selected_shapes): 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 @@ def newShape(self, value=True): 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 @@ def importDirImages(self, dirpath, isDelete=False): 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 @@ def _saveFile(self, annotationFilePath, mode='Manual'): 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 @@ def singleRerecognition(self): 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 @@ def modelChoose(self): 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 @@ def loadFilestate(self, saveDir): 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 @@ def exportJSON(self): 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 @@ def exportJSON(self): # '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 @@ def exportJSON(self): # 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 @@ 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 @@ def initUI(self): 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 @@ def initUI(self): 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 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 @@ def exportJSON(self): 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 @@ def exportJSON(self): 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 @@ def exportJSON(self): # 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("")