프로젝트

일반

사용자정보

통계
| 개정판:

hytos / DTI_PID / DTI_PID / OcrResultDialog.py @ 94c4cd85

이력 | 보기 | 이력해설 | 다운로드 (8.59 KB)

1
# coding: utf-8
2
"""
3
    This is ocr result dialog module
4
"""
5
from PIL import Image
6
import io
7
import numpy as np
8
import math
9
from PyQt5.QtCore import *
10
from PyQt5.QtGui import *
11
from PyQt5.QtWidgets import *
12
import OcrResultDialog_UI
13
import QtImageViewer
14
import tesseract_ocr_module as TOCR
15
from App import App
16
from AppDocData import *
17

    
18
class QOcrResultDialog(QDialog):
19
    def __init__(self, parent, qimage, boundingBox, isModify = False, text = None):
20
        QDialog.__init__(self, parent)
21
        self.textInfoList = None
22

    
23
        self.isModify = isModify
24
        self.image = qimage
25
        self.originImageWidth = qimage.width()
26
        self.originImageHeight = qimage.height()
27
        self.boundingBox = boundingBox
28

    
29
        self.angle = 0 # angle is degree
30

    
31
        self.ui = OcrResultDialog_UI.Ui_Dialog()
32
        self.ui.setupUi(self)
33
        
34
        self.imgW = qimage.width()
35
        self.imgH = qimage.height()
36
        self.image = self.image.scaled(self.imgW, self.imgH)
37
        self.graphicsView = QtImageViewer.QtImageViewer(App.mainWnd())
38
        self.graphicsView.useDefaultCommand() ##### USE DEFAULT COMMAND
39
        self.graphicsView.setImage(self.image)
40
        self.ui.horizontalLayoutGraphicsView.addWidget(self.graphicsView)
41
        
42
        self.ui.counterClockPushButton_2.clicked.connect(lambda : self.rotateImage(True))
43
        self.ui.clockPushButton_2.clicked.connect(lambda : self.rotateImage(False))
44
        self.ui.redetectPushButton_2.clicked.connect(self.detectText)
45
        self.ui.pushButtonMakeTrainingImage.clicked.connect(self.pushButtonMakeTrainingImageClicked)
46

    
47
        if self.isModify == False:
48
            self.detectText()
49
        else:
50
            self.ui.detectResultTextEdit.setPlainText(text)
51

    
52
        self.isAccepted = False
53

    
54
    '''
55
        @brief      Make OCR Training Image
56
        @author     euisung
57
        @date       2018.10.16
58
        @history    euisung     2018.11.02       add notice push
59
    '''
60
    def pushButtonMakeTrainingImageClicked(self):
61
        import uuid
62
        uid = str(uuid.uuid4()) + '.png'
63
        appDocData = AppDocData.instance()
64
        project = appDocData.getCurrentProject()
65
        trainingImgPath = os.path.join(project.getTrainingFilePath(), uid)
66

    
67
        self.image.save(trainingImgPath)
68
        QMessageBox.about(self, "알림", '저장하였습니다.')
69
        
70
    def rotateImage(self, isCounterClock):
71
        for item in self.graphicsView.scene.items():
72
            self.graphicsView.scene.removeItem(item)
73
        self.graphicsView.clearImage()
74
        transform = QTransform()
75
        if isCounterClock:
76
            '''CounterClock'''
77
            self.angle = (self.angle - 90) % 360
78
            transform.rotate(-90)
79
        else:
80
            '''Clock'''
81
            self.angle = (self.angle - 270) % 360
82
            transform.rotate(90)
83
        print(str(360 - self.angle))
84
        self.image = self.image.transformed(transform)
85
        self.graphicsView.setImage(self.image)
86

    
87
    '''
88
        @history 2018.04.26 Jeongwoo    Add Rectangle with modified Coords
89
                 2018.06.20 Jeongwoo    Remove test code
90
                 2018.11.08 euisung     add white char list check process on db
91
                 2018.11.22 euisung     OCR lang apply fixed
92
    '''
93
    def detectText(self):
94
        try:
95
            buffer = QBuffer()
96
            buffer.open(QBuffer.ReadWrite)
97
            self.image.save(buffer, "PNG")
98
            pyImage = Image.open(io.BytesIO(buffer.data()))
99
            img = np.array(pyImage)
100
            
101
            #self.image.save('c:\\temp\\a.png')
102

    
103
            docData = AppDocData.instance()
104
            
105
            # get ocr data of area which has the text
106
            pt = self.boundingBox.center()
107
            areas = [area for area in docData.getAreaList() if area.contains((pt.x(), pt.y()))]
108
            ocr_data = sorted(areas, key=lambda attr: attr.area)[0].OCRData if areas else 'eng'
109
            # up to here
110

    
111
            whiteCharList = docData.getConfigs('Text Recognition', 'White Character List')
112
            if len(whiteCharList) is 0:
113
                self.textInfoList = TOCR.getTextInfo(img, (round(self.boundingBox.x()), round(self.boundingBox.y())), 0, language=ocr_data)
114
            else:
115
                self.textInfoList = TOCR.getTextInfo(img, (round(self.boundingBox.x()), round(self.boundingBox.y())), 0, language=ocr_data, conf = whiteCharList[0].value)
116

    
117
            if self.textInfoList is not None and len(self.textInfoList) > 0:
118
                self.ui.detectResultTextEdit.setText(self.getPlainText(self.textInfoList))
119
                for textInfo in self.textInfoList:
120
                    self.graphicsView.scene.addRect(textInfo.getX()-int(self.boundingBox.x()), textInfo.getY()-int(self.boundingBox.y()), textInfo.getW(), textInfo.getH(), QPen(Qt.red, 1, Qt.SolidLine))
121
            else:
122
                self.ui.detectResultTextEdit.setText("Not Found")
123
        except Exception as ex:
124
            from App import App 
125
            message = 'error occured({}) in {}:{}'.format(ex, sys.exc_info()[-1].tb_frame.f_code.co_filename, sys.exc_info()[-1].tb_lineno)
126
            App.mainWnd().addMessage.emit(MessageType.Error, message)
127

    
128
    def getPlainText(self, textInfoList):
129
        text = ''
130
        for index in range(len(textInfoList)):
131
            textInfo = textInfoList[index]
132
            if index != 0:
133
                text = text + '\n'
134
            text = text + textInfo.getText()
135
        return text
136

    
137
    '''
138
        @brief      OK Button Clicked. Remake TextInfo object
139
        @author     Jeongwoo
140
        @date       18.04.19
141
        @history    18.04.20    Jeongwoo    Calculate Start Point Coordinates by rotated angle
142
                    18.04.26    Jeongwoo    Scene.itemAt(textX - boundBox.x(), textY - boundBox.y())
143
    '''
144
    def accept(self):
145
        self.isAccepted = True
146

    
147
        try:
148
            text = self.ui.detectResultTextEdit.toPlainText()
149
            splitText = text.split('\n')
150
            if self.textInfoList is not None and len(self.textInfoList) > 0:
151
                for index in range(len(splitText)):
152
                    textInfo = self.textInfoList[index]
153
                    item = self.graphicsView.scene.itemAt(QPointF(float(textInfo.getX() - int(self.boundingBox.x())), float(textInfo.getY()-int(self.boundingBox.y()))), QTransform())
154
                    if item is not None:
155
                        ## Transform rectangle for calculate start point
156
                        imgTransform = QTransform()
157
                        if self.angle == 90 or self.angle == 270:
158
                            imgTransform.translate(self.image.height()*0.5, self.image.width()*0.5)
159
                        elif self.angle == 0 or self.angle == 360:
160
                            imgTransform.translate(self.image.width()*0.5, self.image.height()*0.5)
161
                        imgTransform.rotate(-abs(self.angle))
162
                        imgTransform.translate(-self.image.width()*0.5, -self.image.height()*0.5)
163
                        rect = QRect(textInfo.getX() - int(self.boundingBox.x()), textInfo.getY() - int(self.boundingBox.y()), textInfo.getW(), textInfo.getH())
164
                        rect = imgTransform.mapRect(rect)
165
                        ## up to here
166
                        textInfo.setX(rect.x() + int(self.boundingBox.x()))
167
                        textInfo.setY(rect.y() + int(self.boundingBox.y()))
168
                        textInfo.setText(splitText[index])
169
                        radian = round(math.radians(abs(self.angle)), 2)
170
                        textInfo.setAngle(radian) # 360 degree == 6.28319 radian
171
                        if radian == 1.57 or radian == 4.71:
172
                            width = textInfo.getW()
173
                            height = textInfo.getH()
174
                            textInfo.setW(height) ## SWAP
175
                            textInfo.setH(width) ## SWAP
176
                self.textInfoList = self.textInfoList[:len(splitText)]
177
                QDialog.accept(self)
178
            else:
179
                QMessageBox.about(self.ui.ocrDialogButtonBox, "알림", "텍스트 검출을 하신 후 다시 시도해주세요.")
180
        except Exception as ex:
181
            from App import App 
182
            message = 'error occured({}) in {}:{}'.format(ex, sys.exc_info()[-1].tb_frame.f_code.co_filename, sys.exc_info()[-1].tb_lineno)
183
            App.mainWnd().addMessage.emit(MessageType.Error, message)
184

    
185
    def reject(self):
186
        self.isAccepted = False
187
        self.textInfoList = None
188
        QDialog.reject(self)
189

    
190
    '''
191
        @brief  Display this QDialog
192
    '''
193
    def showDialog(self):
194
        #self.setWindowFlags(self.windowFlags() & ~Qt.WindowContextHelpButtonHint)
195
        self.exec_()
196
        return (self.isAccepted, self.textInfoList)