hytos / DTI_PID / DTI_PID / OcrResultDialog.py @ acc6b10d
이력 | 보기 | 이력해설 | 다운로드 (13.8 KB)
| 1 |
# coding: utf-8
|
|---|---|
| 2 |
"""
|
| 3 |
This is ocr result dialog module
|
| 4 |
"""
|
| 5 |
from PIL import Image |
| 6 |
import io, sys |
| 7 |
import numpy as np |
| 8 |
import math |
| 9 |
|
| 10 |
from PyQt5.QtCore import * |
| 11 |
from PyQt5.QtGui import * |
| 12 |
from PyQt5.QtWidgets import * |
| 13 |
import OcrResultDialog_UI |
| 14 |
import QtImageViewer |
| 15 |
import tesseract_ocr_module as TOCR |
| 16 |
from App import App |
| 17 |
from AppDocData import * |
| 18 |
|
| 19 |
|
| 20 |
class SpellTextEdit(QTextEdit): |
| 21 |
def __init__(self, *args): |
| 22 |
QTextEdit.__init__(self, *args)
|
| 23 |
|
| 24 |
# Default dictionary based on the current locale.
|
| 25 |
app_doc_data = AppDocData.instance() |
| 26 |
white_char_list = app_doc_data.getConfigs('Text Recognition', 'White Character List') |
| 27 |
self.highlighter = Highlighter(self.document()) |
| 28 |
self.highlighter.white_char_list = white_char_list[0].value if white_char_list else None |
| 29 |
|
| 30 |
|
| 31 |
class Highlighter(QSyntaxHighlighter): |
| 32 |
err_format = QTextCharFormat() |
| 33 |
err_format.setUnderlineColor(Qt.red) |
| 34 |
err_format.setUnderlineStyle(QTextCharFormat.SpellCheckUnderline) |
| 35 |
|
| 36 |
def __init__(self, *args): |
| 37 |
QSyntaxHighlighter.__init__(self, *args)
|
| 38 |
|
| 39 |
self.white_char_list = None |
| 40 |
|
| 41 |
def highlightBlock(self, text): |
| 42 |
pos = 0
|
| 43 |
for word in text.split(): |
| 44 |
if self.white_char_list and any((c not in self.white_char_list) for c in word): |
| 45 |
self.setFormat(pos, len(word), self.err_format) |
| 46 |
pos += len(word) + 1 |
| 47 |
|
| 48 |
|
| 49 |
class QOcrResultDialog(QDialog): |
| 50 |
def __init__(self, parent, qimage, boundingBox, text_item=None): |
| 51 |
QDialog.__init__(self, parent)
|
| 52 |
self.textInfoList = []
|
| 53 |
|
| 54 |
self._text_item = text_item
|
| 55 |
self.image = qimage
|
| 56 |
self.boundingBox = boundingBox
|
| 57 |
|
| 58 |
self.angle = 0 # angle is degree |
| 59 |
|
| 60 |
self.ui = OcrResultDialog_UI.Ui_Dialog()
|
| 61 |
self.ui.setupUi(self) |
| 62 |
self.ui.detectResultTextEdit = SpellTextEdit()
|
| 63 |
self.ui.detectResultTextEdit.setFont(QFont('Consolas', 15, QFont.Bold)) |
| 64 |
self.ui.horizontalLayoutTextEdit.addWidget(self.ui.detectResultTextEdit) |
| 65 |
|
| 66 |
app_doc_data = AppDocData.instance() |
| 67 |
configs = app_doc_data.getAppConfigs('app', 'mode') |
| 68 |
if configs and 1 == len(configs) and 'advanced' == configs[0].value: |
| 69 |
pass
|
| 70 |
else:
|
| 71 |
self.ui.pushButtonMakeTrainingImage.setVisible(False) |
| 72 |
|
| 73 |
self.graphicsView = QtImageViewer.QtImageViewer(App.mainWnd())
|
| 74 |
self.graphicsView.useDefaultCommand() # USE DEFAULT COMMAND |
| 75 |
self.graphicsView.setImage(self.image) |
| 76 |
self.ui.horizontalLayoutGraphicsView.addWidget(self.graphicsView) |
| 77 |
|
| 78 |
self.ui.counterClockPushButton_2.clicked.connect(lambda: self.rotateImage(True)) |
| 79 |
self.ui.clockPushButton_2.clicked.connect(lambda: self.rotateImage(False)) |
| 80 |
self.ui.redetectPushButton_2.clicked.connect(self.detect_text) |
| 81 |
self.ui.pushButtonMakeTrainingImage.clicked.connect(self.pushButtonMakeTrainingImageClicked) |
| 82 |
|
| 83 |
self.ui.comboBoxOCRData.addItem('eng') |
| 84 |
tessdata_path = os.path.join(os.getenv('ALLUSERSPROFILE'), 'Digital PID', 'Tesseract-OCR', 'tessdata') |
| 85 |
if os.path.isfile(os.path.join(tessdata_path, app_doc_data.getCurrentProject().getName() + '.traineddata')): |
| 86 |
self.ui.comboBoxOCRData.addItem(app_doc_data.getCurrentProject().getName())
|
| 87 |
|
| 88 |
configs = app_doc_data.getConfigs('Text Recognition', 'OCR Data') |
| 89 |
value = configs[0].value if 1 == len(configs) else '' |
| 90 |
if value:
|
| 91 |
at = self.ui.comboBoxOCRData.findText(value)
|
| 92 |
self.ui.comboBoxOCRData.setCurrentIndex(at)
|
| 93 |
else:
|
| 94 |
self.ui.comboBoxOCRData.selectedIndex = 0 |
| 95 |
|
| 96 |
if not self._text_item: |
| 97 |
self.detect_text()
|
| 98 |
else:
|
| 99 |
self.ui.detectResultTextEdit.setPlainText(self._text_item.text()) |
| 100 |
self.ui.checkBoxSeperate.setChecked(False) |
| 101 |
|
| 102 |
self.isAccepted = False |
| 103 |
|
| 104 |
def showEvent(self, QShowEvent): |
| 105 |
"""show event"""
|
| 106 |
self.graphicsView.zoomImageInit()
|
| 107 |
|
| 108 |
'''
|
| 109 |
@brief Make OCR Training Image
|
| 110 |
@author euisung
|
| 111 |
@date 2018.10.16
|
| 112 |
@history euisung 2018.11.02 add notice push
|
| 113 |
'''
|
| 114 |
|
| 115 |
def pushButtonMakeTrainingImageClicked(self): |
| 116 |
import uuid |
| 117 |
uid = str(uuid.uuid4()) + '.png' |
| 118 |
appDocData = AppDocData.instance() |
| 119 |
project = appDocData.getCurrentProject() |
| 120 |
trainingImgPath = os.path.join(project.getTrainingFilePath(), uid) |
| 121 |
|
| 122 |
self.image.save(trainingImgPath)
|
| 123 |
QMessageBox.about(self, self.tr("INFO"), self.tr('Successfully saved.')) |
| 124 |
QDialog.reject(self)
|
| 125 |
|
| 126 |
def rotateImage(self, isCounterClock): |
| 127 |
for item in self.graphicsView.scene.items(): |
| 128 |
self.graphicsView.scene.removeItem(item)
|
| 129 |
self.graphicsView.clearImage()
|
| 130 |
transform = QTransform() |
| 131 |
if isCounterClock:
|
| 132 |
'''CounterClock'''
|
| 133 |
self.angle = (self.angle - 90) % 360 |
| 134 |
transform.rotate(-90)
|
| 135 |
else:
|
| 136 |
'''Clock'''
|
| 137 |
self.angle = (self.angle - 270) % 360 |
| 138 |
transform.rotate(90)
|
| 139 |
# print(str(360 - self.angle))
|
| 140 |
self.image = self.image.transformed(transform) |
| 141 |
self.graphicsView.setImage(self.image) |
| 142 |
self.textInfoList = []
|
| 143 |
|
| 144 |
'''
|
| 145 |
@history 2018.04.26 Jeongwoo Add Rectangle with modified Coords
|
| 146 |
2018.06.20 Jeongwoo Remove test code
|
| 147 |
2018.11.08 euisung add white char list check process on db
|
| 148 |
2018.11.22 euisung OCR lang apply fixed
|
| 149 |
'''
|
| 150 |
def detect_text(self): |
| 151 |
try:
|
| 152 |
buffer = QBuffer() |
| 153 |
buffer.open(QBuffer.ReadWrite)
|
| 154 |
self.image.save(buffer, "PNG") |
| 155 |
pyImage = Image.open(io.BytesIO(buffer.data()))
|
| 156 |
img = np.array(pyImage) |
| 157 |
|
| 158 |
app_doc_data = AppDocData.instance() |
| 159 |
|
| 160 |
ocr_data = self.ui.comboBoxOCRData.currentText()
|
| 161 |
|
| 162 |
whiteCharList = app_doc_data.getConfigs('Text Recognition', 'White Character List') |
| 163 |
if len(whiteCharList) is 0: |
| 164 |
self.textInfoList = TOCR.getTextInfo(img, (round(self.boundingBox.x()), round(self.boundingBox.y())), 0, |
| 165 |
language=ocr_data) |
| 166 |
else:
|
| 167 |
self.textInfoList = TOCR.getTextInfo(img, (round(self.boundingBox.x()), round(self.boundingBox.y())), 0, |
| 168 |
language=ocr_data, conf=whiteCharList[0].value)
|
| 169 |
|
| 170 |
if self.textInfoList is not None and len(self.textInfoList) > 0: |
| 171 |
self.ui.detectResultTextEdit.setText(self.getPlainText(self.textInfoList)) |
| 172 |
for textInfo in self.textInfoList: |
| 173 |
self.graphicsView.scene.addRect(textInfo.getX() - int(self.boundingBox.x()), |
| 174 |
textInfo.getY() - int(self.boundingBox.y()), textInfo.getW(), |
| 175 |
textInfo.getH(), QPen(Qt.red, 1, Qt.SolidLine))
|
| 176 |
else:
|
| 177 |
self.ui.detectResultTextEdit.setText("Not Found") |
| 178 |
except Exception as ex: |
| 179 |
from App import App |
| 180 |
message = 'error occurred({}) in {}:{}'.format(ex, sys.exc_info()[-1].tb_frame.f_code.co_filename, |
| 181 |
sys.exc_info()[-1].tb_lineno)
|
| 182 |
App.mainWnd().addMessage.emit(MessageType.Error, message) |
| 183 |
|
| 184 |
def getPlainText(self, textInfoList): |
| 185 |
text = ''
|
| 186 |
for index in range(len(textInfoList)): |
| 187 |
textInfo = textInfoList[index] |
| 188 |
if index != 0: |
| 189 |
text = text + '\n'
|
| 190 |
text = text + textInfo.getText() |
| 191 |
return text
|
| 192 |
|
| 193 |
'''
|
| 194 |
@brief OK Button Clicked. Remake TextInfo object
|
| 195 |
@author Jeongwoo
|
| 196 |
@date 18.04.19
|
| 197 |
@history 18.04.20 Jeongwoo Calculate Start Point Coordinates by rotated angle
|
| 198 |
18.04.26 Jeongwoo Scene.itemAt(textX - boundBox.x(), textY - boundBox.y())
|
| 199 |
'''
|
| 200 |
|
| 201 |
def accept(self): |
| 202 |
from TextInfo import TextInfo |
| 203 |
self.isAccepted = True |
| 204 |
|
| 205 |
try:
|
| 206 |
text = self.ui.detectResultTextEdit.toPlainText()
|
| 207 |
if text == '' or text == 'Not Found': |
| 208 |
QMessageBox.about(self.ui.ocrDialogButtonBox, self.tr('Notice'), |
| 209 |
self.tr('Please try again after recognition or type.')) |
| 210 |
return
|
| 211 |
|
| 212 |
isSplit = self.ui.checkBoxSeperate.isChecked()
|
| 213 |
if isSplit:
|
| 214 |
splitText = text.split('\n')
|
| 215 |
else:
|
| 216 |
splitText = [text] |
| 217 |
|
| 218 |
# try to detect text if there is no result of detection or
|
| 219 |
# count of text info list not match with count of split text
|
| 220 |
if isSplit and self.textInfoList and (len(self.textInfoList) == len(splitText)): |
| 221 |
for index in range(len(self.textInfoList)): |
| 222 |
self.textInfoList[index].setText(splitText[index])
|
| 223 |
elif isSplit and not self.textInfoList: |
| 224 |
self.detect_text()
|
| 225 |
if len(self.textInfoList) == len(splitText[index]): |
| 226 |
for index in range(len(self.textInfoList)): |
| 227 |
self.textInfoList[index].setText(splitText[index])
|
| 228 |
else:
|
| 229 |
self.textInfoList = self.getMergedTextInfo(text) |
| 230 |
else:
|
| 231 |
self.textInfoList = self.getMergedTextInfo(text) |
| 232 |
|
| 233 |
'''
|
| 234 |
minX, minY, maxX, maxY = sys.maxsize, sys.maxsize, 0, 0
|
| 235 |
for textInfo in self.textInfoList:
|
| 236 |
x, y, w, h = textInfo.getX(), textInfo.getY(), textInfo.getW(), textInfo.getH()
|
| 237 |
minX = min(x, minX)
|
| 238 |
minY = min(y, minY)
|
| 239 |
maxX = max(x + w, maxX)
|
| 240 |
maxY = max(y + h, maxY)
|
| 241 |
|
| 242 |
self.textInfoList = [TextInfo(text, minX, minY, maxX - minX, maxY - minY, 0)]
|
| 243 |
'''
|
| 244 |
|
| 245 |
'''
|
| 246 |
if self.textInfoList:
|
| 247 |
if self._text_item:
|
| 248 |
textInfo = self.textInfoList[0]
|
| 249 |
# Transform rectangle for calculate start point
|
| 250 |
imgTransform = QTransform()
|
| 251 |
if self.angle == 90 or self.angle == 270:
|
| 252 |
imgTransform.translate(self.image.height() * 0.5, self.image.width() * 0.5)
|
| 253 |
elif self.angle == 0 or self.angle == 360:
|
| 254 |
imgTransform.translate(self.image.width() * 0.5, self.image.height() * 0.5)
|
| 255 |
imgTransform.rotate(-abs(self.angle))
|
| 256 |
imgTransform.translate(-self.image.width() * 0.5, -self.image.height() * 0.5)
|
| 257 |
rect = QRect(textInfo.getX() - int(self.boundingBox.x()),
|
| 258 |
textInfo.getY() - int(self.boundingBox.y()), textInfo.getW(), textInfo.getH())
|
| 259 |
rect = imgTransform.mapRect(rect)
|
| 260 |
# up to here
|
| 261 |
textInfo.setX(rect.x() + int(self.boundingBox.x()))
|
| 262 |
textInfo.setY(rect.y() + int(self.boundingBox.y()))
|
| 263 |
textInfo.setText(splitText[0])
|
| 264 |
radian = round(math.radians(abs(self.angle)), 2)
|
| 265 |
textInfo.setAngle(radian) # 360 degree == 6.28319 radian
|
| 266 |
if radian == 1.57 or radian == 4.71:
|
| 267 |
width = textInfo.getW()
|
| 268 |
height = textInfo.getH()
|
| 269 |
textInfo.setW(height) # SWAP
|
| 270 |
textInfo.setH(width) # SWAP
|
| 271 |
|
| 272 |
self._text_item.setPlainText(splitText[0])
|
| 273 |
self._text_item.loc = [textInfo.getX(), textInfo.getY()]
|
| 274 |
self._text_item.size = (textInfo.getW(), textInfo.getH())
|
| 275 |
self._text_item.angle = radian
|
| 276 |
self._text_item.update_shape()
|
| 277 |
|
| 278 |
self.textInfoList = self.textInfoList[1:]
|
| 279 |
'''
|
| 280 |
|
| 281 |
QDialog.accept(self)
|
| 282 |
|
| 283 |
except Exception as ex: |
| 284 |
from App import App |
| 285 |
message = 'error occurred({}) in {}:{}'.format(ex, sys.exc_info()[-1].tb_frame.f_code.co_filename, |
| 286 |
sys.exc_info()[-1].tb_lineno)
|
| 287 |
App.mainWnd().addMessage.emit(MessageType.Error, message) |
| 288 |
|
| 289 |
def getMergedTextInfo(self, text): |
| 290 |
import cv2 |
| 291 |
from TextInfo import TextInfo |
| 292 |
|
| 293 |
buffer = QBuffer() |
| 294 |
buffer.open(QBuffer.ReadWrite)
|
| 295 |
self.image.save(buffer, "PNG") |
| 296 |
pyImage = Image.open(io.BytesIO(buffer.data()))
|
| 297 |
img = np.array(pyImage) |
| 298 |
|
| 299 |
img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) |
| 300 |
imgNot = np.ones(img.shape, np.uint8) |
| 301 |
cv2.bitwise_not(img, imgNot) |
| 302 |
imgNot = cv2.dilate(imgNot, np.ones((8, 8), np.uint8)) |
| 303 |
|
| 304 |
contours, hierarchy = cv2.findContours(imgNot, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) |
| 305 |
minX, minY, maxX, maxY = sys.maxsize, sys.maxsize, 0, 0 |
| 306 |
if len(contours) is 0: |
| 307 |
minX, minY, maxX, maxY = self.boundingBox.x(), self.boundingBox.y(), self.boundingBox.x() + self.image.width(), self.boundingBox.y() + self.image.height() |
| 308 |
else:
|
| 309 |
minX, minY, maxX, maxY = sys.maxsize, sys.maxsize, 0, 0 |
| 310 |
for cnt in contours: |
| 311 |
x, y, w, h = cv2.boundingRect(cnt) |
| 312 |
minX = min(x, minX)
|
| 313 |
minY = min(y, minY)
|
| 314 |
maxX = max(x + w, maxX)
|
| 315 |
maxY = max(y + h, maxY)
|
| 316 |
minX, minY, maxX, maxY = minX + self.boundingBox.x(), minY + self.boundingBox.y(), maxX + self.boundingBox.x(), maxY + self.boundingBox.y() |
| 317 |
|
| 318 |
return [TextInfo(text, minX, minY, maxX - minX, maxY - minY, 0)] |
| 319 |
|
| 320 |
def reject(self): |
| 321 |
self.isAccepted = False |
| 322 |
self.textInfoList = None |
| 323 |
QDialog.reject(self)
|
| 324 |
|
| 325 |
'''
|
| 326 |
@brief Display this QDialog
|
| 327 |
'''
|
| 328 |
|
| 329 |
def showDialog(self): |
| 330 |
# self.setWindowFlags(self.windowFlags() & ~Qt.WindowContextHelpButtonHint)
|
| 331 |
self.exec_()
|
| 332 |
return (self.isAccepted, self.textInfoList) |
